From a73cced38ac113b30924fb4cefdbbbece91c68d8 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:23:05 -0700 Subject: [PATCH 01/19] perf: scale agent execution with typed execution state, a claim dispatcher, and owned continuation Replace history-derived scheduling with an agent execution engine. Each agent has a typed execution head of pointers, monotonic boundaries and summaries recomputed at commit, plus immutable opening lineage on model contexts, config changes and checkpoints. One pure selector decides next work for claim, owned continuation, release and recovery. Semantic transitions run through a cell-bound unit of work that plans locks up front in sorted order, supports savepoints, and publishes each changed agent's head and queue in two pipelined statements before enqueueing webhooks and post-commit notifications. Execution writes live in a private query package enforced by SQL rules and a Go analyzer. Workers run a bounded claim dispatcher, keep an owned agent across ready steps within a configurable budget, and accept model output, select and prepare the next call in one transaction. Migration 052 builds every existing agent's execution state from history, validates it, and drops the history frontier functions. It requires a writer-stop maintenance window. Also fixes: - continuations stranded when a turn carried unanswered inputs from an earlier turn - child-before-parent and machine/parent lock-order hazards - subagent launch reusing a machine grant deleted while it waited - project and organization deletion skipping event webhooks Co-Authored-By: Claude Opus 5.5 (1M context) --- .golangci.yml | 9 +- Makefile | 4 +- cmd/worker/main.go | 5 + compose.yaml | 1 + docs/self-hosting/configuration.mdx | 11 + docs/self-hosting/deployment.mdx | 7 + internal/agentmessage/message.go | 33 + internal/compaction/runner.go | 13 +- internal/compaction/runner_candidate_test.go | 20 +- internal/config/config.go | 27 + internal/config/config_test.go | 29 + internal/config/worker_continuation_test.go | 48 + .../execution_cutover_integration_test.go | 70 + ...tegration_owned_schema_integration_test.go | 1 + .../dbmigrate/postgres_integration_test.go | 44 +- ...ck_integration_cutover_integration_test.go | 493 ++++-- ...tegration_pending_work_integration_test.go | 242 ++- .../e2e/service_e2e_deterministic_test.go | 66 +- .../e2e/service_e2e_durable_history_test.go | 22 +- internal/harness/kernel/agent_executor.go | 40 +- ...or_compaction_boundary_integration_test.go | 138 +- ...compaction_progressive_integration_test.go | 53 +- ...or_compaction_provider_integration_test.go | 234 ++- .../harness/kernel/agent_executor_context.go | 70 +- .../kernel/agent_executor_model_context.go | 9 + ...executor_model_context_integration_test.go | 19 +- ...executor_model_request_integration_test.go | 16 +- ...utor_output_truncation_integration_test.go | 114 +- .../harness/kernel/agent_executor_response.go | 33 +- .../agent_executor_retry_integration_test.go | 157 +- ...retry_runtime_recovery_integration_test.go | 40 +- .../agent_executor_test_helpers_test.go | 131 +- ...xecutor_tool_lifecycle_integration_test.go | 6 + .../kernel/agent_executor_tool_work.go | 150 +- ...ent_executor_web_tools_integration_test.go | 16 +- internal/harness/kernel/types.go | 7 + internal/harness/tools/approvals.go | 59 +- internal/harness/tools/dispatch.go | 7 +- .../dispatch_transaction_integration_test.go | 75 +- internal/harness/tools/executor.go | 49 +- .../tools/model_context_test_helpers_test.go | 21 +- internal/harness/worker/continuation.go | 54 + .../worker/continuation_integration_test.go | 210 +++ internal/harness/worker/continuation_test.go | 266 +++ .../worker/round_trips_integration_test.go | 106 ++ internal/harness/worker/worker.go | 337 ++-- .../harness/worker/worker_integration_test.go | 131 +- internal/harness/worker/worker_test.go | 183 +- .../agent_interactions_integration_test.go | 46 +- .../agents_list_route_integration_test.go | 19 +- internal/httpapi/discord_app_setup.go | 7 - internal/httpapi/github_app_setup.go | 7 - .../model_context_test_helpers_test.go | 17 +- internal/httpapi/options_integration_test.go | 23 + internal/httpapi/server.go | 18 +- .../slack_event_routes_integration_test.go | 6 +- internal/interactionform/answers.go | 47 + internal/interactionform/form_test.go | 27 + internal/interactionform/question.go | 26 + .../machinepool/manager_integration_test.go | 420 +++-- internal/modelcontext/context.go | 7 + internal/modelcontext/prepared.go | 52 + internal/modelcontext/prepared_test.go | 69 + internal/modelcontext/types.go | 1 + internal/modelprotocol/retry.go | 30 + internal/modelretry/policy.go | 3 +- internal/notifications/notifications.go | 35 + .../notifications_snapshot_test.go | 34 + internal/processresult/results.go | 299 ++++ .../results_test.go} | 12 +- .../storage/artifactstore/artifacts_store.go | 2 +- .../artifactstore/prepared_artifacts.go | 2 +- ...default_reconciliation_integration_test.go | 79 +- .../storage/executionstore/actors_store.go | 2 +- ...nt_config_change_gates_integration_test.go | 91 +- .../agent_config_changes_store.go | 174 +- .../executionstore/agent_configs_store.go | 21 +- .../agent_content_inputs_store.go | 119 +- .../agent_fixture_integration.go | 20 +- .../agent_input_admission_integration_test.go | 37 +- ...agent_input_admission_test_helpers_test.go | 73 +- .../executionstore/agent_input_rows.go | 158 +- .../agent_inputs_backlog_integration_test.go | 19 +- .../agent_inputs_backlog_store.go | 108 +- .../agent_inputs_ingress_store.go | 54 +- .../executionstore/agent_inputs_store.go | 673 +------ .../agent_integration_targets.go | 18 +- .../agent_interaction_cancellation_store.go | 80 - .../agent_interactions_store.go | 766 +------- .../agent_interactions_store_test.go | 125 -- .../executionstore/agent_launch_inbox.go | 17 +- .../executionstore/agent_launch_input.go | 9 +- .../agent_launch_integration_test.go | 491 +++-- .../executionstore/agent_launch_machines.go | 6 +- .../executionstore/agent_launch_store.go | 60 +- .../agent_machine_source_reconciliation.go | 5 +- internal/storage/executionstore/agent_rows.go | 25 - .../agent_runtime_cancellation_store.go | 422 +---- .../agent_runtime_lock_reaping_store.go | 113 +- .../agent_runtime_locks_integration_test.go | 121 +- .../agent_runtime_locks_store.go | 404 +---- .../agent_runtime_recovery_store.go | 159 -- .../executionstore/agent_runtime_store.go | 240 ++- .../executionstore/agent_turns_store.go | 138 -- .../agent_work_integration_test.go | 511 ++++++ .../executionstore/agent_work_store.go | 51 + .../executionstore/agents_integration_test.go | 346 +++- .../context_checkpoints_boundary.go | 71 - .../context_checkpoints_store.go | 210 +-- .../context_checkpoints_terminal_failure.go | 89 +- .../executionstore/context_events_store.go | 6 +- .../cron_trigger_agent_inputs_store.go | 15 +- internal/storage/executionstore/db_values.go | 7 - .../executionstore/default_reconciliation.go | 4 +- .../event_webhooks_integration_test.go | 241 ++- .../executionstore/event_webhooks_store.go | 76 - .../executionstore/execution_content.go | 26 + .../executionstore/execution_models.go | 65 + .../executionstore/execution_mutation.go | 30 + ...cution_unit_lock_order_integration_test.go | 329 ++++ ...xecution_unit_webhooks_integration_test.go | 84 + .../storage/executionstore/execution_work.go | 153 ++ .../executionstore/export_integration_test.go | 548 ++++-- .../inbox_input_integration_test.go | 175 +- .../executionstore/inbox_input_store.go | 20 +- .../executionstore/inbox_subscription.go | 3 +- .../integration_config_activation.go | 2 +- .../integration_handler_targets.go | 4 +- .../integration_input_origin.go | 17 +- ...tegration_interactions_integration_test.go | 73 +- .../integration_interactions_store.go | 282 ++- .../integration_testmain_test.go | 8 - .../interaction_selection_admission.go | 43 +- ...ction_selection_export_integration_test.go | 39 +- ...on_selection_producers_integration_test.go | 27 +- ...tion_selection_queries_integration_test.go | 140 +- .../storage/executionstore/json_helpers.go | 12 - ..._authority_checkpoints_integration_test.go | 81 +- ..._authority_constraints_integration_test.go | 208 +-- ...authority_continuation_integration_test.go | 375 ++-- ...vent_authority_helpers_integration_test.go | 223 +-- ..._authority_model_calls_integration_test.go | 72 +- .../kernel_event_authority_model_outputs.go | 60 +- .../kernel_event_authority_store.go | 649 ------- .../machine_daemon_bootstrap_store.go | 35 +- .../machine_daemon_reconciliation_store.go | 102 +- .../machine_daemon_runtime_expiry_store.go | 21 +- .../machine_daemon_runtime_store.go | 37 +- .../executionstore/machine_daemon_store.go | 37 +- .../executionstore/machine_lifecycle_store.go | 50 +- .../machine_lifecycle_terminal_work.go | 398 ++--- .../executionstore/machine_pools_store.go | 33 +- ...del_call_schema_guards_integration_test.go | 110 +- .../executionstore/model_context_data.go | 83 + .../model_context_data_integration_test.go | 43 + .../executionstore/model_contexts_boundary.go | 86 - .../executionstore/model_contexts_claims.go | 578 +----- .../model_contexts_compaction.go | 362 +--- .../model_contexts_lifecycle.go | 246 +-- .../model_contexts_lifecycle_test.go | 80 - .../model_contexts_preparation.go | 45 + ...l_contexts_preparation_integration_test.go | 140 ++ .../model_contexts_replay_integration_test.go | 16 +- .../executionstore/model_contexts_store.go | 39 - .../model_outputs_content_projection.go | 219 --- .../model_outputs_terminal_error.go | 308 +--- .../executionstore/model_outputs_workflow.go | 697 +------- .../model_outputs_workflow_test.go | 229 --- .../storage/executionstore/model_usage.go | 36 - .../executionstore/model_usage_test.go | 57 - .../executionstore/model_work_transition.go | 57 + .../model_work_transition_integration_test.go | 214 +++ .../executionstore/org_lifecycle_ops.go | 15 +- .../pool_machines_integration_test.go | 254 ++- .../process_action_completion_store.go | 216 +-- .../executionstore/process_actions_store.go | 33 +- ...s_cancellation_archive_integration_test.go | 144 +- ...action_tool_completion_integration_test.go | 91 +- .../process_daemon_completion_store.go | 403 +---- ...l_ready_frontier_races_integration_test.go | 48 +- .../executionstore/process_read_results.go | 252 +-- .../executionstore/process_result_facts.go | 21 + ...cess_retry_model_ready_integration_test.go | 112 +- .../process_runtime_integration_test.go | 239 ++- ...runtime_lease_recovery_integration_test.go | 248 ++- ...s_start_cancel_integration_helpers_test.go | 143 +- .../process_tool_unreachable_store.go | 227 +-- ...ss_unanswered_openings_integration_test.go | 117 +- ...wakeup_input_admission_integration_test.go | 276 +-- .../storage/executionstore/processes_store.go | 24 +- .../processes_test_helpers_test.go | 141 +- .../project_machine_pool_grants_store.go | 17 +- .../storage/executionstore/provider_cost.go | 8 - .../executionstore/provider_cost_test.go | 6 +- .../provider_runtime_reconciliation_store.go | 35 +- internal/storage/executionstore/store.go | 79 +- .../subagents_integration_test.go | 4 +- .../storage/executionstore/subagents_store.go | 339 ++-- .../executionstore/tool_call_commands.go | 10 +- .../executionstore/tool_call_execution.go | 132 +- .../executionstore/tool_call_reads_store.go | 2 +- .../executionstore/tool_calls_store.go | 593 ++----- .../tool_calls_test_helpers_test.go | 18 +- .../executionstore/tool_calls_workflow.go | 220 --- .../tool_calls_workflow_test.go | 116 +- .../tool_overflow_integration_test.go | 129 +- .../executionstore/tool_result_overflow.go | 141 +- .../tool_result_overflow_test.go | 56 +- internal/storage/executionstore/tool_rows.go | 152 -- .../executionstore/tool_work_preparation.go | 52 + internal/storage/executionstore/turn_rows.go | 38 - .../identitystore/principal_activity.go | 4 +- .../inbox_integration_test.go | 4 +- .../storage/integrationstore/inbox_lease.go | 4 +- .../storage/integrationstore/inbox_store.go | 2 +- .../inbox_terminal_integration_test.go | 2 +- .../integration_agent_conversation.go | 4 +- .../integration_credential_authority.go | 4 +- .../integration_launch_claims.go | 2 +- ...egration_launch_claims_integration_test.go | 2 +- .../integrationstore/integration_lifecycle.go | 12 +- .../integrationstore/integration_locks.go | 5 +- ...ration_profile_choices_integration_test.go | 7 +- .../integrationstore/integration_routing.go | 3 +- .../integrationstore/integration_runtime.go | 2 +- .../integrationstore/integration_setup.go | 2 +- .../integration_subscription_registration.go | 6 +- ...ation_target_lifecycle_integration_test.go | 74 +- .../integration_target_lookup.go | 2 +- .../integrationstore/integration_targets.go | 4 +- .../integration_targets_integration_test.go | 2 +- .../storage/integrationstore/integrations.go | 2 +- .../integrations_integration_test.go | 7 +- .../integrationstore/scheduled_event.go | 2 +- internal/storage/integrationstore/store.go | 16 +- .../storage/internal/agentexecution/cell.go | 63 + .../agentexecution/commands_agents.go | 254 +++ .../commands_backlog_integration_test.go | 145 ++ .../agentexecution/commands_cancel.go | 238 +++ .../agentexecution/commands_children.go | 127 ++ .../agentexecution/commands_compaction.go | 347 ++++ .../agentexecution/commands_config.go | 143 ++ .../agentexecution/commands_content.go | 224 +++ .../commands_custody_integration_test.go | 217 +++ .../agentexecution/commands_failure.go | 265 +++ .../commands_family_integration_test.go | 317 ++++ .../agentexecution/commands_inputs.go | 325 ++++ .../commands_integration_test.go | 959 ++++++++++ .../agentexecution/commands_interactions.go | 459 +++++ .../agentexecution/commands_lifecycle.go | 305 ++++ .../agentexecution/commands_models.go | 321 ++++ .../agentexecution/commands_outputs.go | 369 ++++ .../commands_permission_integration_test.go | 144 ++ .../agentexecution/commands_processes.go | 305 ++++ .../commands_processes_integration_test.go | 374 ++++ .../agentexecution/commands_runtime.go | 422 +++++ .../commands_runtime_integration_test.go | 228 +++ .../commands_selection_integration_test.go | 111 ++ .../internal/agentexecution/commands_tools.go | 339 ++++ .../internal/agentexecution/database.go | 138 ++ .../internal/agentexecution/evidence_test.go | 171 ++ .../fixture_integration_test.go | 82 + .../storage/internal/agentexecution/handle.go | 47 + .../internal/agentexecution/history_test.go | 306 ++++ .../internal/executiondb/batch.go | 548 ++++++ .../executiondb/commands_agents.sql.go | 181 ++ .../executiondb/commands_events.sql.go | 249 +++ .../executiondb/commands_inputs.sql.go | 242 +++ .../executiondb/commands_interactions.sql.go | 202 +++ .../executiondb/commands_models.sql.go | 349 ++++ .../executiondb/commands_outputs.sql.go | 258 +++ .../executiondb/commands_processes.sql.go | 108 ++ .../executiondb/commands_runtime.sql.go | 348 ++++ .../executiondb/commands_tools.sql.go | 222 +++ .../internal/executiondb/contexts.sql.go | 114 ++ .../internal/executiondb/db_gen.go | 26 + .../internal/executiondb/inputs.sql.go | 107 ++ .../internal/executiondb/load.sql.go | 81 + .../internal/executiondb/models_gen.go | 5 + .../internal/executiondb/outputs.sql.go | 62 + .../internal/executiondb/publication.sql.go | 6 + .../internal/executiondb/state.sql.go | 58 + .../internal/executiondb/tools.sql.go | 73 + .../internal/executiondb/turns_events.sql.go | 79 + .../lineage_integration_test.go | 133 ++ .../linearity_integration_test.go | 133 ++ .../storage/internal/agentexecution/load.go | 190 ++ .../load_count_integration_test.go | 65 + .../agentexecution/load_integration_test.go | 48 + .../storage/internal/agentexecution/locks.go | 254 +++ .../internal/agentexecution/mutation.go | 159 ++ .../payload_plans_integration_test.go | 32 + .../publication_integration_test.go | 144 ++ .../queries/commands_agents.sql | 36 + .../queries/commands_events.sql | 45 + .../queries/commands_inputs.sql | 40 + .../queries/commands_interactions.sql | 34 + .../queries/commands_models.sql | 69 + .../queries/commands_outputs.sql | 40 + .../queries/commands_processes.sql | 10 + .../queries/commands_runtime.sql | 68 + .../agentexecution/queries/commands_tools.sql | 40 + .../agentexecution/queries/contexts.sql | 35 + .../agentexecution/queries/inputs.sql | 30 + .../internal/agentexecution/queries/load.sql | 271 +++ .../agentexecution/queries/outputs.sql | 5 + .../agentexecution/queries/publication.sql | 43 + .../internal/agentexecution/queries/state.sql | 13 + .../internal/agentexecution/queries/tools.sql | 13 + .../agentexecution/queries/turns_events.sql | 16 + .../reconstruction_integration_test.go | 536 ++++++ .../storage/internal/agentexecution/repair.go | 216 +++ .../storage/internal/agentexecution/route.go | 23 + .../agentexecution/schema_integration_test.go | 317 ++++ .../internal/agentexecution/selection.go | 283 +++ .../selection_invariants_test.go | 191 ++ .../internal/agentexecution/selection_load.go | 198 +++ .../internal/agentexecution/selection_test.go | 415 +++++ .../internal/agentexecution/snapshot.go | 63 + .../agentexecution/sqlc_integration_test.go | 26 + .../storage/internal/agentexecution/state.go | 198 +++ .../storage/internal/agentexecution/unit.go | 189 ++ .../agentexecution/unit_integration_test.go | 272 +++ .../internal/agentexecution/webhooks.go | 112 ++ .../internal/dbsqlc/agent_inputs.sql.go | 813 --------- .../dbsqlc/agent_inputs_ingress.sql.go | 91 +- .../internal/dbsqlc/agent_interactions.sql.go | 443 ----- .../dbsqlc/agent_machine_bindings.sql.go | 26 - .../internal/dbsqlc/agent_runtime.sql.go | 142 -- .../dbsqlc/agent_runtime_locks.sql.go | 368 ---- .../internal/dbsqlc/agent_turns.sql.go | 163 +- .../storage/internal/dbsqlc/agent_work.sql.go | 131 -- .../storage/internal/dbsqlc/agents.sql.go | 22 - .../dbsqlc/context_checkpoints.sql.go | 110 -- .../storage/internal/dbsqlc/events.sql.go | 81 - .../dbsqlc/integration_interactions.sql.go | 76 - .../integration_target_lifecycle.sql.go | 21 - .../dbsqlc/kernel_event_authority.sql.go | 953 +--------- .../internal/dbsqlc/model_contexts.sql.go | 794 --------- .../storage/internal/dbsqlc/processes.sql.go | 133 +- .../storage/internal/dbsqlc/secrets.sql.go | 63 + .../storage/internal/dbsqlc/subagents.sql.go | 12 +- .../storage/internal/dbsqlc/tool_calls.sql.go | 1577 +---------------- .../internal/lifecyclelock/lifecyclelock.go | 38 +- .../storage/internal/secretops/secretops.go | 3 +- .../internal/storeutil/transactions.go | 16 - .../internal/storeutil/transactions_test.go | 93 - .../storage/orglifecycle/orgs_projects.go | 80 +- internal/storage/orglifecycle/service.go | 37 +- internal/storage/queries/actors.sql | 12 +- internal/storage/queries/agent_inputs.sql | 397 ----- .../storage/queries/agent_inputs_ingress.sql | 35 +- .../storage/queries/agent_interactions.sql | 127 -- .../queries/agent_machine_bindings.sql | 15 - internal/storage/queries/agent_runtime.sql | 58 - .../storage/queries/agent_runtime_locks.sql | 130 -- internal/storage/queries/agent_turns.sql | 74 +- internal/storage/queries/agent_work.sql | 49 - internal/storage/queries/agents.sql | 7 - .../storage/queries/context_checkpoints.sql | 48 - internal/storage/queries/events.sql | 20 - internal/storage/queries/identity.sql | 10 +- .../queries/integration_interactions.sql | 29 - .../queries/integration_target_lifecycle.sql | 10 - .../queries/kernel_event_authority.sql | 356 +--- internal/storage/queries/model_contexts.sql | 429 ----- internal/storage/queries/processes.sql | 31 +- internal/storage/queries/secrets.sql | 10 + internal/storage/queries/subagents.sql | 12 +- internal/storage/queries/tool_calls.sql | 648 +------ .../resource_limits_integration_test.go | 28 +- ...secret_payload_current_integration_test.go | 71 + .../secretstore/secret_payload_store.go | 37 +- internal/storage/secretstore/secret_rows.go | 18 + internal/storage/store.go | 10 +- internal/storage/store_integration.go | 11 + internal/storage/storeerr/errors.go | 24 +- migrations/000052_agent_execution_state.sql | 417 +++++ .../000053_agent_payload_query_plans.sql | 30 + sqlc.vet-db.yaml | 48 +- sqlc.yaml | 22 +- tools/ci/sqlrules/checker.go | 146 +- tools/ci/sqlrules/checker_test.go | 109 +- tools/omnaralint/execution.go | 145 ++ tools/omnaralint/go.mod | 6 + tools/omnaralint/go.sum | 24 +- tools/omnaralint/omnaralint.go | 12 +- tools/omnaralint/omnaralint_test.go | 67 + .../github.com/jackc/pgx/v5/pgconn/pgconn.go | 3 + .../src/github.com/jackc/pgx/v5/pgx.go | 8 + .../storage/executionstore/escapes.go | 62 + 391 files changed, 28790 insertions(+), 22339 deletions(-) create mode 100644 internal/agentmessage/message.go create mode 100644 internal/config/worker_continuation_test.go create mode 100644 internal/dbmigrate/execution_cutover_integration_test.go create mode 100644 internal/harness/worker/continuation.go create mode 100644 internal/harness/worker/continuation_integration_test.go create mode 100644 internal/harness/worker/continuation_test.go create mode 100644 internal/harness/worker/round_trips_integration_test.go delete mode 100644 internal/httpapi/discord_app_setup.go delete mode 100644 internal/httpapi/github_app_setup.go create mode 100644 internal/httpapi/options_integration_test.go create mode 100644 internal/interactionform/answers.go create mode 100644 internal/interactionform/question.go create mode 100644 internal/modelcontext/prepared.go create mode 100644 internal/modelcontext/prepared_test.go create mode 100644 internal/modelprotocol/retry.go create mode 100644 internal/notifications/notifications_snapshot_test.go create mode 100644 internal/processresult/results.go rename internal/{storage/executionstore/process_read_results_test.go => processresult/results_test.go} (88%) delete mode 100644 internal/storage/executionstore/agent_interaction_cancellation_store.go delete mode 100644 internal/storage/executionstore/agent_interactions_store_test.go delete mode 100644 internal/storage/executionstore/agent_runtime_recovery_store.go create mode 100644 internal/storage/executionstore/agent_work_integration_test.go create mode 100644 internal/storage/executionstore/agent_work_store.go delete mode 100644 internal/storage/executionstore/context_checkpoints_boundary.go create mode 100644 internal/storage/executionstore/execution_content.go create mode 100644 internal/storage/executionstore/execution_models.go create mode 100644 internal/storage/executionstore/execution_mutation.go create mode 100644 internal/storage/executionstore/execution_unit_lock_order_integration_test.go create mode 100644 internal/storage/executionstore/execution_unit_webhooks_integration_test.go create mode 100644 internal/storage/executionstore/execution_work.go create mode 100644 internal/storage/executionstore/model_context_data.go create mode 100644 internal/storage/executionstore/model_context_data_integration_test.go delete mode 100644 internal/storage/executionstore/model_contexts_boundary.go delete mode 100644 internal/storage/executionstore/model_contexts_lifecycle_test.go create mode 100644 internal/storage/executionstore/model_contexts_preparation.go create mode 100644 internal/storage/executionstore/model_contexts_preparation_integration_test.go delete mode 100644 internal/storage/executionstore/model_outputs_content_projection.go delete mode 100644 internal/storage/executionstore/model_outputs_workflow_test.go delete mode 100644 internal/storage/executionstore/model_usage_test.go create mode 100644 internal/storage/executionstore/model_work_transition.go create mode 100644 internal/storage/executionstore/model_work_transition_integration_test.go create mode 100644 internal/storage/executionstore/process_result_facts.go create mode 100644 internal/storage/executionstore/tool_work_preparation.go create mode 100644 internal/storage/internal/agentexecution/cell.go create mode 100644 internal/storage/internal/agentexecution/commands_agents.go create mode 100644 internal/storage/internal/agentexecution/commands_backlog_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_cancel.go create mode 100644 internal/storage/internal/agentexecution/commands_children.go create mode 100644 internal/storage/internal/agentexecution/commands_compaction.go create mode 100644 internal/storage/internal/agentexecution/commands_config.go create mode 100644 internal/storage/internal/agentexecution/commands_content.go create mode 100644 internal/storage/internal/agentexecution/commands_custody_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_failure.go create mode 100644 internal/storage/internal/agentexecution/commands_family_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_inputs.go create mode 100644 internal/storage/internal/agentexecution/commands_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_interactions.go create mode 100644 internal/storage/internal/agentexecution/commands_lifecycle.go create mode 100644 internal/storage/internal/agentexecution/commands_models.go create mode 100644 internal/storage/internal/agentexecution/commands_outputs.go create mode 100644 internal/storage/internal/agentexecution/commands_permission_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_processes.go create mode 100644 internal/storage/internal/agentexecution/commands_processes_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_runtime.go create mode 100644 internal/storage/internal/agentexecution/commands_runtime_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_selection_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_tools.go create mode 100644 internal/storage/internal/agentexecution/database.go create mode 100644 internal/storage/internal/agentexecution/evidence_test.go create mode 100644 internal/storage/internal/agentexecution/fixture_integration_test.go create mode 100644 internal/storage/internal/agentexecution/handle.go create mode 100644 internal/storage/internal/agentexecution/history_test.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/batch.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_agents.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_events.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_interactions.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_models.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_outputs.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_processes.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_runtime.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/commands_tools.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/contexts.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/db_gen.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/inputs.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/load.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/models_gen.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/outputs.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/publication.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/state.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/tools.sql.go create mode 100644 internal/storage/internal/agentexecution/internal/executiondb/turns_events.sql.go create mode 100644 internal/storage/internal/agentexecution/lineage_integration_test.go create mode 100644 internal/storage/internal/agentexecution/linearity_integration_test.go create mode 100644 internal/storage/internal/agentexecution/load.go create mode 100644 internal/storage/internal/agentexecution/load_count_integration_test.go create mode 100644 internal/storage/internal/agentexecution/load_integration_test.go create mode 100644 internal/storage/internal/agentexecution/locks.go create mode 100644 internal/storage/internal/agentexecution/mutation.go create mode 100644 internal/storage/internal/agentexecution/payload_plans_integration_test.go create mode 100644 internal/storage/internal/agentexecution/publication_integration_test.go create mode 100644 internal/storage/internal/agentexecution/queries/commands_agents.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_events.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_inputs.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_interactions.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_models.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_outputs.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_processes.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_runtime.sql create mode 100644 internal/storage/internal/agentexecution/queries/commands_tools.sql create mode 100644 internal/storage/internal/agentexecution/queries/contexts.sql create mode 100644 internal/storage/internal/agentexecution/queries/inputs.sql create mode 100644 internal/storage/internal/agentexecution/queries/load.sql create mode 100644 internal/storage/internal/agentexecution/queries/outputs.sql create mode 100644 internal/storage/internal/agentexecution/queries/publication.sql create mode 100644 internal/storage/internal/agentexecution/queries/state.sql create mode 100644 internal/storage/internal/agentexecution/queries/tools.sql create mode 100644 internal/storage/internal/agentexecution/queries/turns_events.sql create mode 100644 internal/storage/internal/agentexecution/reconstruction_integration_test.go create mode 100644 internal/storage/internal/agentexecution/repair.go create mode 100644 internal/storage/internal/agentexecution/route.go create mode 100644 internal/storage/internal/agentexecution/schema_integration_test.go create mode 100644 internal/storage/internal/agentexecution/selection.go create mode 100644 internal/storage/internal/agentexecution/selection_invariants_test.go create mode 100644 internal/storage/internal/agentexecution/selection_load.go create mode 100644 internal/storage/internal/agentexecution/selection_test.go create mode 100644 internal/storage/internal/agentexecution/snapshot.go create mode 100644 internal/storage/internal/agentexecution/sqlc_integration_test.go create mode 100644 internal/storage/internal/agentexecution/state.go create mode 100644 internal/storage/internal/agentexecution/unit.go create mode 100644 internal/storage/internal/agentexecution/unit_integration_test.go create mode 100644 internal/storage/internal/agentexecution/webhooks.go delete mode 100644 internal/storage/internal/dbsqlc/agent_work.sql.go delete mode 100644 internal/storage/queries/agent_work.sql create mode 100644 internal/storage/secret_payload_current_integration_test.go create mode 100644 internal/storage/store_integration.go create mode 100644 migrations/000052_agent_execution_state.sql create mode 100644 migrations/000053_agent_payload_query_plans.sql create mode 100644 tools/omnaralint/execution.go create mode 100644 tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgconn/pgconn.go create mode 100644 tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgx.go create mode 100644 tools/omnaralint/testdata/src/github.com/omnara-ai/omnara/internal/storage/executionstore/escapes.go diff --git a/.golangci.yml b/.golangci.yml index 372df6376..054337145 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -141,8 +141,11 @@ linters: exclusions: generated: disable rules: - - path: internal/storage/internal/storeutil/pool\.go - source: '\.Query\(ctx, sql, args\.\.\.\)' + - path: internal/storage/internal/agentexecution/internal/executiondb/batch\.go + linters: + - intrange + - path: internal/storage/internal/(storeutil/pool|agentexecution/database)\.go + source: '(\.Query\(ctx, (sql|query), args\.\.\.\)|b\.BatchResults\.Query\(\))' linters: - rowserrcheck # oapi-codegen's oneOf wrappers intentionally mix value accessors with @@ -180,7 +183,7 @@ linters: - linters: - lll source: '^\s*\)\.[A-Za-z]' - - path: (internal/storage/internal/dbsqlc/|internal/machinedaemon/statedb/internal/dbsqlc/) + - path: (internal/storage/internal/dbsqlc/|internal/storage/internal/agentexecution/internal/executiondb/|internal/machinedaemon/statedb/internal/dbsqlc/) linters: - errcheck - govet diff --git a/Makefile b/Makefile index 0f91a871a..7463a495b 100644 --- a/Makefile +++ b/Makefile @@ -23,6 +23,7 @@ MIGRATION_CHECK ?= $(GO) -C tools/ci run ./migrationcheck OMNARAD_VERSION ?= 0.0.0-dev MIGRATION_DIRS := migrations internal/machinedaemon/statedb/migrations SQLC_OWNED_PATHS := internal/storage/internal/dbsqlc internal/storage/queries \ + internal/storage/internal/agentexecution/internal/executiondb internal/storage/internal/agentexecution/queries \ internal/machinedaemon/statedb/internal/dbsqlc internal/machinedaemon/statedb/queries OMNARALINT_SOURCES := $(shell find tools/omnaralint -name '*.go' -print) GO_MODULE_DIRS := . observability tools/ci tools/goose tools/omnaralint @@ -259,6 +260,7 @@ sqlc-check: sql-rules: CGO_ENABLED=0 $(GO) -C tools/ci run ./sqlrules ../../internal/storage/queries + CGO_ENABLED=0 $(GO) -C tools/ci run ./sqlrules ../../internal/storage/internal/agentexecution/queries sqlc-vet: $(SQLC) vet -f sqlc.yaml @@ -267,7 +269,7 @@ migrate-test-db: $(TEST_DB_ENV) $(GO) test -count=1 -run '^TestMigrateLocalDatabase$$' ./internal/testutil/integrationdb sqlc-vet-db: - SQLC_DATABASE_URL=$${SQLC_DATABASE_URL:-$${OMNARA_TEST_DATABASE_URL:-$(TEST_DATABASE_URL)}} $(SQLC) vet -f sqlc.vet-db.yaml + OMNARA_TEST_DATABASE_URL=$${SQLC_DATABASE_URL:-$${OMNARA_TEST_DATABASE_URL:-$(TEST_DATABASE_URL)}} $(GO) test -count=1 -tags=integration ./internal/storage/internal/agentexecution -run '^TestExecutionSQLCPrepare$$' sqlc-vet-local-db: db-up migrate-test-db sqlc-vet-db diff --git a/cmd/worker/main.go b/cmd/worker/main.go index 62fa8d46d..ae08a90ee 100644 --- a/cmd/worker/main.go +++ b/cmd/worker/main.go @@ -241,8 +241,13 @@ func main() { StreamLog: log, }) kernelWorker := workerpkg.NewWorker(store.Execution(), executor, workerpkg.Options{ + Continuation: &workerpkg.ContinuationOptions{ + MaxModelStarts: cfg.WorkerContinuationMaxModelStarts, + MaxDuration: cfg.WorkerContinuationMaxDuration, + }, Log: log, Capacity: cfg.WorkerCapacity, + ClaimConcurrency: cfg.WorkerClaimConcurrency, AsyncToolCapacity: cfg.WorkerAsyncToolCapacity, ControlSubscriber: redisBus, }) diff --git a/compose.yaml b/compose.yaml index 39f5fd094..409af6eda 100644 --- a/compose.yaml +++ b/compose.yaml @@ -126,6 +126,7 @@ services: OMNARA_WORKER_EVENT_WEBHOOK_CONCURRENCY: ${OMNARA_WORKER_EVENT_WEBHOOK_CONCURRENCY:-128} OMNARA_EVENT_WEBHOOK_PER_ORG_CONCURRENCY: ${OMNARA_EVENT_WEBHOOK_PER_ORG_CONCURRENCY:-128} OMNARA_WORKER_CAPACITY: ${OMNARA_WORKER_CAPACITY:-4} + OMNARA_WORKER_CLAIM_CONCURRENCY: ${OMNARA_WORKER_CLAIM_CONCURRENCY:-4} OMNARA_WORKER_INBOX_CAPACITY: ${OMNARA_WORKER_INBOX_CAPACITY:-4} OMNARA_WORKER_DISCORD_CAPACITY: ${OMNARA_WORKER_DISCORD_CAPACITY:-1024} OMNARA_WORKER_ASYNC_TOOL_CAPACITY: ${OMNARA_WORKER_ASYNC_TOOL_CAPACITY:-32} diff --git a/docs/self-hosting/configuration.mdx b/docs/self-hosting/configuration.mdx index d206b9cf1..2a27e1577 100644 --- a/docs/self-hosting/configuration.mdx +++ b/docs/self-hosting/configuration.mdx @@ -28,11 +28,22 @@ All configuration is via environment variables. Use `.env.example` in the repo a | `OMNARA_WORKER_EVENT_WEBHOOK_CONCURRENCY` | Maximum event webhook deliveries in flight per worker process | `128` | No | | `OMNARA_EVENT_WEBHOOK_PER_ORG_CONCURRENCY` | Soft cap on concurrent event webhook deliveries per org across workers; use the same value on every worker | `128` | No | | `OMNARA_WORKER_CAPACITY` | Maximum agent turns executed concurrently by each worker process | `4` | No | +| `OMNARA_WORKER_CLAIM_CONCURRENCY` | Maximum claim transactions each worker process runs at once; keep it small, since claims share one queue | `4` | No | +| `OMNARA_WORKER_CONTINUATION_MAX_MODEL_STARTS` | Additional model calls a worker may start for an agent it already owns before releasing it back to the queue; `0` releases after every step | `2` | No | +| `OMNARA_WORKER_CONTINUATION_MAX_DURATION` | Longest a worker keeps continuing an owned agent, checked between steps; never interrupts a running step (Go duration) | `30s` | No | | `OMNARA_WORKER_INBOX_CAPACITY` | Maximum integration inbox receipts processed concurrently per worker process; integer from 1 through 100 | `4` | No | | `OMNARA_WORKER_ASYNC_TOOL_CAPACITY` | Maximum runtime-owned asynchronous tool calls in flight per worker process | `32` | No | | `OMNARA_WORKER_BACKGROUND_TOOL_CAPACITY` | Background tool concurrency and in-memory queue capacity per worker process | `8` | No | | `OMNARA_MAINTENANCE_INTERVAL` | Maintenance loop interval (Go duration) | `1s` | No | +### Agent execution + +Agents have typed execution state maintained by the execution store. Each transition +updates domain records and publishes the agent's execution state and wakeup in the +same transaction. Workers claim wakeups under fenced runtime leases and may +continue ready work within the configured continuation budget. Rollbacks discard +all pending state and notifications. + ### Integration inbox operation Inbox concurrency is independent of agent-turn concurrency. All worker activity diff --git a/docs/self-hosting/deployment.mdx b/docs/self-hosting/deployment.mdx index bf572a913..1717c4539 100644 --- a/docs/self-hosting/deployment.mdx +++ b/docs/self-hosting/deployment.mdx @@ -25,6 +25,13 @@ Open `http://localhost:8000` and sign up — the local email driver is `console` ## Production database upgrades +The execution-state migration requires a writer-stop maintenance window. Stop API, +worker, maintenance, and daemon-report handlers before running migrations, and +restart them together on the new release after migration succeeds. Existing agent +history is backfilled before the migration completes. Do not allow an older binary +to resume writing, delete runtime leases, or clear external tool ownership during +the upgrade. Unfinished work recovers through the normal lease and retry paths. + Run the release's one-shot `migrate` image before starting its API, worker, and maintenance services. Connect the migrator directly to the PostgreSQL writer; do not use RDS Proxy or transaction-pooling PgBouncer. Long-running services diff --git a/internal/agentmessage/message.go b/internal/agentmessage/message.go new file mode 100644 index 000000000..4fae951e0 --- /dev/null +++ b/internal/agentmessage/message.go @@ -0,0 +1,33 @@ +package agentmessage + +import ( + "fmt" + "strings" +) + +func Render(displayName, key, childPublicID, kind, text string) string { + label := fmt.Sprintf( + "Subagent %q (agent_id %s, key %q)", displayName, childPublicID, key, + ) + var header string + switch kind { + case "result": + header = label + " finished its turn:" + case "refused": + header = label + " stopped because its model refused to continue; it produced no final answer:" + case "content_filtered": + header = label + " stopped because its model output was blocked by a content filter; " + + "it produced no final answer:" + case "failed": + header = label + " failed:" + case "question": + header = label + " asked a question and is paused until a human answers it. " + + "Messaging it with send_agent_message cancels the question." + default: + header = label + ":" + } + if strings.TrimSpace(text) == "" { + return header + } + return header + "\n\n" + text +} diff --git a/internal/compaction/runner.go b/internal/compaction/runner.go index b6608143e..9348f8f79 100644 --- a/internal/compaction/runner.go +++ b/internal/compaction/runner.go @@ -131,7 +131,8 @@ func (r Runner) run( } resolved, err := r.Resolver.Resolve(ctx, selection) if err != nil { - if ctx.Err() != nil && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) { + if ctx.Err() != nil && + (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) { return RunResult{}, err } return r.recordPreSendFailure( @@ -265,7 +266,8 @@ func (r Runner) run( providerAttempt.ProviderRequestStarted = true providerAttempt.Response = response if err != nil { - if ctx.Err() != nil && (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) { + if ctx.Err() != nil && + (errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) { return RunResult{}, err } if _, classified := model.ClassifyError(err); !classified { @@ -401,7 +403,7 @@ func (r Runner) run( }, ) if err != nil { - if errors.Is(err, executionstore.ErrCheckpointBoundaryUnsafe) { + if errors.Is(err, storeerr.ErrCheckpointBoundaryUnsafe) { nextEnd, adjustmentErr := r.nextSmallerSourceEnd(ctx, input.Plan) if adjustmentErr != nil { return RunResult{}, errors.Join(err, adjustmentErr) @@ -509,7 +511,10 @@ func (r Runner) validate(input RunInput) error { return nil } -func (r Runner) resultForUnclaimed(ctx context.Context, claim executionstore.ModelCallClaim) (RunResult, error) { +func (r Runner) resultForUnclaimed( + ctx context.Context, + claim executionstore.ModelCallClaim, +) (RunResult, error) { result := RunResult{ ModelCallContextID: claim.Context.ID, } diff --git a/internal/compaction/runner_candidate_test.go b/internal/compaction/runner_candidate_test.go index a55b60d27..692d840fa 100644 --- a/internal/compaction/runner_candidate_test.go +++ b/internal/compaction/runner_candidate_test.go @@ -9,6 +9,7 @@ import ( "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/model" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/stretchr/testify/require" ) @@ -114,7 +115,7 @@ func TestRunnerReplansWhenPublishFindsUnsafeBoundary(t *testing.T) { textCompactionEvent(1, strings.Repeat("first closed event ", 20)), textCompactionEvent(2, strings.Repeat("second closed event ", 20)), }, - publishErrs: []error{executionstore.ErrCheckpointBoundaryUnsafe}, + publishErrs: []error{storeerr.ErrCheckpointBoundaryUnsafe}, } client := &summaryModel{results: []summaryResult{ {response: completeSummaryResponse("Concise first summary.")}, @@ -307,8 +308,10 @@ func TestSafeCompactionSourceEndsRequireModelVisibleSemantics(t *testing.T) { } func TestRenderEventSourceReducesProjectionWithoutMutatingCanonicalOutput(t *testing.T) { - original := json.RawMessage(`[{"type":"structured_data","value":{"stdout":"` + strings.Repeat("x", 40_960) + - `TAIL"}}]`) + original := json.RawMessage( + `[{"type":"structured_data","value":{"stdout":"` + strings.Repeat("x", 40_960) + + `TAIL"}}]`, + ) event := mustCompactionEvent(1, "tool_result", "completed", append(json.RawMessage(nil), original...)) event.ToolName = "run_command" event.ProviderCallID = "call_1" @@ -446,8 +449,10 @@ func TestCompactionAllowanceFitPrefersWholeTurnAndPreservesPartialProgress(t *te }{ {name: "whole turn fits reduced allowance", firstText: "short request", secondText: largeText, firstKind: string(events.KindAgentInput), wantEnd: 2, wantReduced: true}, - {name: "oversized turn keeps preferred partial", firstText: "short request", secondText: largeText + largeText, - firstKind: string(events.KindAgentInput), wantEnd: 1}, + {name: "oversized turn keeps preferred partial", + firstText: "short request", + secondText: largeText + largeText, + firstKind: string(events.KindAgentInput), wantEnd: 1}, {name: "oversized turn keeps reduced partial", firstText: largeText, secondText: largeText, firstKind: string(events.KindModelOutput), wantEnd: 1, wantReduced: true}, {name: "preferred whole turn wins", firstText: "small output", secondText: largeText, @@ -497,7 +502,10 @@ func (m summaryModelWithExplicitCapabilities) Capabilities() model.Capabilities func TestCompactionRetainsOverflowLocator(t *testing.T) { path := "/artifacts/art_overflow" parts, err := json.Marshal([]map[string]any{ - {"type": "structured_data", "value": map[string]any{"path": path, "preview": strings.Repeat("x", 8192)}}, + { + "type": "structured_data", + "value": map[string]any{"path": path, "preview": strings.Repeat("x", 8192)}, + }, {"type": "text", "text": strings.Repeat("other content", 5000)}, }) if err != nil { diff --git a/internal/config/config.go b/internal/config/config.go index 53356297a..4767a6810 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -66,7 +66,10 @@ type Config struct { AllowInsecureDev bool WorkerEventWebhookConcurrency int EventWebhookPerOrgConcurrency int + WorkerContinuationMaxModelStarts int + WorkerContinuationMaxDuration time.Duration WorkerCapacity int + WorkerClaimConcurrency int WorkerInboxCapacity int WorkerDiscordCapacity int WorkerAsyncToolCapacity int @@ -214,10 +217,22 @@ func Load() (Config, error) { if err != nil { return Config{}, err } + workerContinuationMaxModelStarts, err := getenvInt("OMNARA_WORKER_CONTINUATION_MAX_MODEL_STARTS", 2) + if err != nil { + return Config{}, err + } + workerContinuationMaxDuration, err := time.ParseDuration(getenv("OMNARA_WORKER_CONTINUATION_MAX_DURATION", "30s")) + if err != nil { + return Config{}, fmt.Errorf("parse OMNARA_WORKER_CONTINUATION_MAX_DURATION: %w", err) + } workerCapacity, err := getenvInt("OMNARA_WORKER_CAPACITY", 4) if err != nil { return Config{}, err } + workerClaimConcurrency, err := getenvInt("OMNARA_WORKER_CLAIM_CONCURRENCY", 4) + if err != nil { + return Config{}, err + } workerInboxCapacity, err := getenvInt("OMNARA_WORKER_INBOX_CAPACITY", 4) if err != nil { return Config{}, err @@ -266,7 +281,10 @@ func Load() (Config, error) { AllowInsecureDev: os.Getenv("OMNARA_ALLOW_INSECURE_DEV_DEFAULTS") == "1", WorkerEventWebhookConcurrency: workerEventWebhookConcurrency, EventWebhookPerOrgConcurrency: eventWebhookPerOrgConcurrency, + WorkerContinuationMaxModelStarts: workerContinuationMaxModelStarts, + WorkerContinuationMaxDuration: workerContinuationMaxDuration, WorkerCapacity: workerCapacity, + WorkerClaimConcurrency: workerClaimConcurrency, WorkerInboxCapacity: workerInboxCapacity, WorkerDiscordCapacity: workerDiscordCapacity, WorkerAsyncToolCapacity: workerAsyncToolCapacity, @@ -622,9 +640,18 @@ func (cfg Config) ValidateWorker() error { if cfg.EventWebhookPerOrgConcurrency <= 0 { return fmt.Errorf("OMNARA_EVENT_WEBHOOK_PER_ORG_CONCURRENCY must be positive") } + if cfg.WorkerContinuationMaxModelStarts < 0 { + return fmt.Errorf("OMNARA_WORKER_CONTINUATION_MAX_MODEL_STARTS must be non-negative") + } + if cfg.WorkerContinuationMaxDuration <= 0 { + return fmt.Errorf("OMNARA_WORKER_CONTINUATION_MAX_DURATION must be positive") + } if cfg.WorkerCapacity <= 0 { return fmt.Errorf("OMNARA_WORKER_CAPACITY must be positive") } + if cfg.WorkerClaimConcurrency <= 0 { + return fmt.Errorf("OMNARA_WORKER_CLAIM_CONCURRENCY must be positive") + } if cfg.WorkerInboxCapacity < 1 || cfg.WorkerInboxCapacity > 100 { return fmt.Errorf("OMNARA_WORKER_INBOX_CAPACITY must be between 1 and 100") } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 8be0f58f1..45fb85f83 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1106,6 +1106,35 @@ func TestValidateWorkerLoadsExplicitWorkerCapacity(t *testing.T) { } } +func TestValidateWorkerLoadsExplicitClaimConcurrency(t *testing.T) { + t.Setenv("OMNARA_ALLOW_INSECURE_DEV_DEFAULTS", "1") + t.Setenv("OMNARA_WORKER_CLAIM_CONCURRENCY", "2") + + cfg, err := Load() + if err != nil { + t.Fatalf("load config: %v", err) + } + if err := cfg.ValidateWorker(); err != nil { + t.Fatalf("validate worker: %v", err) + } + if cfg.WorkerClaimConcurrency != 2 { + t.Fatalf("unexpected claim concurrency: %+v", cfg) + } +} + +func TestValidateWorkerRejectsInvalidClaimConcurrency(t *testing.T) { + t.Setenv("OMNARA_ALLOW_INSECURE_DEV_DEFAULTS", "1") + t.Setenv("OMNARA_WORKER_CLAIM_CONCURRENCY", "0") + + cfg, err := Load() + if err != nil { + t.Fatalf("load config: %v", err) + } + if err := cfg.ValidateWorker(); err == nil { + t.Fatal("expected invalid claim concurrency error") + } +} + func TestValidateWorkerLoadsExplicitAsyncToolCapacity(t *testing.T) { t.Setenv("OMNARA_ALLOW_INSECURE_DEV_DEFAULTS", "1") t.Setenv("OMNARA_WORKER_ASYNC_TOOL_CAPACITY", "64") diff --git a/internal/config/worker_continuation_test.go b/internal/config/worker_continuation_test.go new file mode 100644 index 000000000..c0f9861ec --- /dev/null +++ b/internal/config/worker_continuation_test.go @@ -0,0 +1,48 @@ +package config + +import ( + "testing" + "time" +) + +func TestWorkerContinuationConfig(t *testing.T) { + for _, test := range []struct { + name, starts, duration string + wantStarts int + wantDuration time.Duration + invalid bool + }{ + {"defaults", "", "", 2, 30 * time.Second, false}, + {"explicit", "3", "12s", 3, 12 * time.Second, false}, + {"disabled", "0", "30s", 0, 30 * time.Second, false}, + {"negative starts", "-1", "30s", 0, 0, true}, + {"invalid starts", "bad", "30s", 0, 0, true}, + {"zero duration", "2", "0s", 0, 0, true}, + {"negative duration", "2", "-1s", 0, 0, true}, + {"invalid duration", "2", "bad", 0, 0, true}, + } { + t.Run(test.name, func(t *testing.T) { + t.Setenv("OMNARA_ALLOW_INSECURE_DEV_DEFAULTS", "1") + t.Setenv("OMNARA_WORKER_CONTINUATION_MAX_MODEL_STARTS", test.starts) + t.Setenv("OMNARA_WORKER_CONTINUATION_MAX_DURATION", test.duration) + cfg, err := Load() + if err == nil { + err = cfg.ValidateWorker() + } + if test.invalid { + if err == nil { + t.Fatal("expected invalid continuation config") + } + return + } + if err != nil { + t.Fatal(err) + } + if cfg.WorkerContinuationMaxModelStarts != test.wantStarts || + cfg.WorkerContinuationMaxDuration != test.wantDuration { + t.Fatalf("continuation=%d/%s want=%d/%s", + cfg.WorkerContinuationMaxModelStarts, cfg.WorkerContinuationMaxDuration, test.wantStarts, test.wantDuration) + } + }) + } +} diff --git a/internal/dbmigrate/execution_cutover_integration_test.go b/internal/dbmigrate/execution_cutover_integration_test.go new file mode 100644 index 000000000..da2f83888 --- /dev/null +++ b/internal/dbmigrate/execution_cutover_integration_test.go @@ -0,0 +1,70 @@ +//go:build integration + +package dbmigrate_test + +import ( + "context" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestExecutionCutoverRemovesHistoryEngine(t *testing.T) { + t.Parallel() + ctx := context.Background() + pool, _ := openPostgresMigrationTestDB(t, ctx) + names := []string{ + "agent_scheduler_state", + "model_call_context_turns", + "agent_stop_events", + "agent_turn_id_at_event_sequence", + "agent_turn_opening_content_inputs", + "agent_latest_turn_id", + "agent_model_call_opening_content_inputs", + "agent_latest_unstarted_model_call_opening_inputs", + "model_call_context_has_later_semantic_event", + "agent_continuable_model_contexts", + "agent_unconsumed_context_checkpoint_frontiers", + "agent_unconsumed_config_change_frontiers", + "agent_tool_work_frontiers", + "agent_has_incomplete_tool_batch", + "agent_model_result_frontiers", + "agent_next_model_work", + "agent_next_wakeup_ready_at", + "refresh_agent_scheduler", + "stage_agent_scheduler", + "read_agent_scheduler", + "flush_scheduler_transaction", + "omnara.scheduler", + } + rows, err := pool.Query( + ctx, + `SELECT 'function',p.proname,pg_get_functiondef(p.oid) FROM pg_proc p JOIN pg_namespace n ON n.oid=p.pronamespace WHERE n.nspname='public' AND p.prokind='f' + UNION ALL SELECT 'view',c.relname,pg_get_viewdef(c.oid) FROM pg_class c JOIN pg_namespace + n ON n.oid=c.relnamespace WHERE n.nspname='public' AND c.relkind IN ('v','m') + UNION ALL SELECT 'relation',c.relname,'' FROM pg_class c JOIN pg_namespace n ON + n.oid=c.relnamespace WHERE n.nspname='public' + UNION ALL SELECT 'trigger',t.tgname,pg_get_triggerdef(t.oid) FROM pg_trigger t JOIN + pg_class c ON c.oid=t.tgrelid JOIN pg_namespace n ON n.oid=c.relnamespace WHERE + n.nspname='public' AND NOT t.tgisinternal`, + ) + require.NoError(t, err) + defer rows.Close() + for rows.Next() { + var kind, name, body string + require.NoError(t, rows.Scan(&kind, &name, &body)) + for _, removed := range names { + require.NotEqual(t, removed, name, "%s", kind) + require.False( + t, + strings.Contains(body, removed), + "%s %s still references %s", + kind, + name, + removed, + ) + } + } + require.NoError(t, rows.Err()) +} diff --git a/internal/dbmigrate/integration_owned_schema_integration_test.go b/internal/dbmigrate/integration_owned_schema_integration_test.go index 7b814e9cf..469db55cf 100644 --- a/internal/dbmigrate/integration_owned_schema_integration_test.go +++ b/internal/dbmigrate/integration_owned_schema_integration_test.go @@ -192,6 +192,7 @@ func TestIntegrationActorMigrationPreservesLegacySlackAttribution(t *testing.T) require.NoError(t, pool.QueryRow(ctx, `SELECT i.actor_id FROM agent_events e JOIN agent_inputs i ON i.id=e.agent_input_id WHERE e.id=$1`, eventID).Scan(&historicalActorID)) require.Equal(t, id, historicalActorID, "historical inputs/events retain their actor UUID") + require.NoError(t, applyProductionPostgresMigrations(ctx, db)) for range 2 { actor, err := executionstore.IntegrationActorParams(integrationstore.IntegrationRecord{ ID: uuid.New(), ProjectID: ids.ProjectID, IntegrationKind: integrationdefinition.SlackThread, diff --git a/internal/dbmigrate/postgres_integration_test.go b/internal/dbmigrate/postgres_integration_test.go index 3b7c1e7de..3fe6c68e4 100644 --- a/internal/dbmigrate/postgres_integration_test.go +++ b/internal/dbmigrate/postgres_integration_test.go @@ -103,8 +103,12 @@ INSERT INTO agent_mcp_connections VALUES require.True(t, defaultIsNull) for _, state := range []string{"initializing", "ready", "failed", "expired"} { for _, value := range []sql.NullString{{}, {String: "", Valid: true}, {String: "failure", Valid: true}} { - _, err := db.ExecContext(ctx, - `UPDATE agent_mcp_connections SET state = $1, initialize_error = $2 WHERE id = 9`, state, value) + _, err := db.ExecContext( + ctx, + `UPDATE agent_mcp_connections SET state = $1, initialize_error = $2 WHERE id = 9`, + state, + value, + ) if !value.Valid || (value.String != "" && (state == "failed" || state == "expired")) { require.NoError(t, err, "state=%s error=%+v", state, value) } else { @@ -259,12 +263,18 @@ INSERT INTO projects(org_id, name, created_at, updated_at) VALUES ($1, 'immutable-project', statement_timestamp(), statement_timestamp()) RETURNING id::text `, originalOrgID).Scan(&projectID)) - if _, err := db.ExecContext(ctx, `UPDATE projects SET org_id = $1 WHERE id = $2`, otherOrgID, projectID); err == nil { + if _, + err := db.ExecContext(ctx, + `UPDATE projects SET org_id = $1 WHERE id = $2`, + otherOrgID, + projectID); err == nil { t.Fatal("project organization update succeeded") } var retainedOrgID string require.NoError( - t, db.QueryRowContext(ctx, `SELECT org_id::text FROM projects WHERE id = $1`, projectID).Scan(&retainedOrgID), + t, + db.QueryRowContext(ctx, `SELECT org_id::text FROM projects WHERE id = $1`, projectID). + Scan(&retainedOrgID), ) if retainedOrgID != originalOrgID { t.Fatalf("project org_id = %s, want %s", retainedOrgID, originalOrgID) @@ -355,7 +365,11 @@ func TestPostgresExplicitConstraintNamesStayBelowTruncationLimit(t *testing.T) { name = strings.ReplaceAll(strings.Trim(name, `"`), `""`, `"`) } if len([]byte(name)) >= 63 { - t.Errorf("explicit constraint name %q in %s reaches PostgreSQL's 63-byte truncation limit", name, migrationFile) + t.Errorf( + "explicit constraint name %q in %s reaches PostgreSQL's 63-byte truncation limit", + name, + migrationFile, + ) } } } @@ -442,7 +456,10 @@ SELECT coalesce(string_agg( FROM redundant_unique_constraints `).Scan(&redundant)) if redundant != "" { - t.Fatalf("unique constraints redundantly contain primary keys without a referencing foreign key: %s", redundant) + t.Fatalf( + "unique constraints redundantly contain primary keys without a referencing foreign key: %s", + redundant, + ) } } @@ -545,7 +562,10 @@ func mutatingForeignKeysWithoutSupportingIndexes( ) string { t.Helper() var unsupported string - require.NoError(t, db.QueryRowContext(ctx, mutatingForeignKeysWithoutSupportingIndexesQuery).Scan(&unsupported)) + require.NoError( + t, + db.QueryRowContext(ctx, mutatingForeignKeysWithoutSupportingIndexesQuery).Scan(&unsupported), + ) return unsupported } @@ -805,7 +825,10 @@ WHERE id = $1 `, grantID); err != nil { t.Fatalf("set valid grant minimums: %v", err) } - if _, err := db.ExecContext(ctx, `UPDATE machine_pools SET min_machine_cpu = 9 WHERE id = $1`, poolID); err == nil { + if _, + err := db.ExecContext(ctx, + `UPDATE machine_pools SET min_machine_cpu = 9 WHERE id = $1`, + poolID); err == nil { t.Fatal("pool minimum above maximum unexpectedly succeeded") } if _, err := db.ExecContext( @@ -1269,7 +1292,10 @@ func TestProcessExecutionSpecMigration(t *testing.T) { {1, "echo hello", "default", "pipe"}, {2, "python", "bash", "pty"}, } { var spec processcmd.ExecutionSpec - require.NoError(t, pool.QueryRow(ctx, `SELECT execution_spec FROM processes WHERE id = $1`, expected.id).Scan(&spec)) + require.NoError( + t, + pool.QueryRow(ctx, `SELECT execution_spec FROM processes WHERE id = $1`, expected.id).Scan(&spec), + ) require.Equal(t, processcmd.ForShell( expected.command, processcmd.ShellSelector(expected.shell), processcmd.IOMode(expected.ioMode), ), spec) diff --git a/internal/dbmigrate/slack_integration_cutover_integration_test.go b/internal/dbmigrate/slack_integration_cutover_integration_test.go index 8242d36d1..99e894382 100644 --- a/internal/dbmigrate/slack_integration_cutover_integration_test.go +++ b/internal/dbmigrate/slack_integration_cutover_integration_test.go @@ -54,11 +54,43 @@ func TestSlackIntegrationCutoverPreflight(t *testing.T) { func TestSlackIntegrationCutoverPreservesConversationsAndHistory(t *testing.T) { for _, scenario := range []string{ - "normal", "memory_store", "pre_internal_ids", "injected_failure", "rewrite_failure", "continuable_retry", "custom_collision", + "normal", + "memory_store", + "pre_internal_ids", + "injected_failure", + "rewrite_failure", + "continuable_retry", + "custom_collision", "live_lease", "unsupported_setup", - "fixed_agent", "mislabeled_fixed_agent", "multiple_targets", "multiple_integrations", "enabled", "json", "yaml", "yaml_alias", "enabled_json", "enabled_yaml", "channel", "dm", - "wire_json", "wire_yaml", "enabled_wire_json", "enabled_wire_yaml", "enabled_null_json", "enabled_null_yaml", - "absent_send", "preflight_invalid_policy", "preflight_unmapped_policy", "preflight_address", "preflight_quota", "invalid_policy", "unmapped_policy", "invalid_address", "bad_hash", "bad_source_hash", "config_limit", + "fixed_agent", + "mislabeled_fixed_agent", + "multiple_targets", + "multiple_integrations", + "enabled", + "json", + "yaml", + "yaml_alias", + "enabled_json", + "enabled_yaml", + "channel", + "dm", + "wire_json", + "wire_yaml", + "enabled_wire_json", + "enabled_wire_yaml", + "enabled_null_json", + "enabled_null_yaml", + "absent_send", + "preflight_invalid_policy", + "preflight_unmapped_policy", + "preflight_address", + "preflight_quota", + "invalid_policy", + "unmapped_policy", + "invalid_address", + "bad_hash", + "bad_source_hash", + "config_limit", } { t.Run(scenario, func(t *testing.T) { ctx := t.Context() @@ -112,7 +144,10 @@ func TestSlackIntegrationCutoverPreservesConversationsAndHistory(t *testing.T) { var compiledObject map[string]any require.NoError(t, json.Unmarshal([]byte(compiled), &compiledObject)) legacyTools := testutil.RequireType[map[string]any](t, compiledObject["tools"]) - legacySendPolicy := testutil.RequireType[map[string]any](t, legacyTools["send_integration_message"]) + legacySendPolicy := testutil.RequireType[map[string]any]( + t, + legacyTools["send_integration_message"], + ) sendingEnabled := strings.HasPrefix(scenario, "enabled") || scenario == "absent_send" if sendingEnabled { legacySendPolicy["enabled"] = true @@ -138,8 +173,14 @@ func TestSlackIntegrationCutoverPreservesConversationsAndHistory(t *testing.T) { } memoryStoreID := uuid.New() if scenario == "memory_store" { - exec(`INSERT INTO memory_stores(id,project_id,name) VALUES($1,$2,'review-notes')`, memoryStoreID, ids.ProjectID) - compiledObject["memory_stores"] = []map[string]any{{"id": memoryStoreID, "access": "read_write"}} + exec( + `INSERT INTO memory_stores(id,project_id,name) VALUES($1,$2,'review-notes')`, + memoryStoreID, + ids.ProjectID, + ) + compiledObject["memory_stores"] = []map[string]any{ + {"id": memoryStoreID, "access": "read_write"}, + } } canonical, err := json.Marshal(compiledObject) require.NoError(t, err) @@ -185,7 +226,11 @@ tools: } } if source != nil && sendingEnabled { - source = strings.ReplaceAll(testutil.RequireType[string](t, source), "enabled: false", "enabled: true") + source = strings.ReplaceAll( + testutil.RequireType[string](t, source), + "enabled: false", + "enabled: true", + ) } if strings.Contains(scenario, "wire_") || strings.Contains(scenario, "null_") { policy := `{"type":"built_in","enabled":false}` @@ -204,8 +249,13 @@ tools: } if source != nil { exec(`ALTER TABLE agent_configs DISABLE TRIGGER agent_configs_immutable`) - exec(`UPDATE agent_configs SET source=$2,source_format=$3,source_hash=$4 WHERE id=$1`, - configID, source, sourceFormat, fmt.Sprintf("%x", sha256.Sum256([]byte(testutil.RequireType[string](t, source))))) + exec( + `UPDATE agent_configs SET source=$2,source_format=$3,source_hash=$4 WHERE id=$1`, + configID, + source, + sourceFormat, + fmt.Sprintf("%x", sha256.Sum256([]byte(testutil.RequireType[string](t, source)))), + ) exec(`ALTER TABLE agent_configs ENABLE TRIGGER agent_configs_immutable`) } @@ -228,13 +278,20 @@ tools: ids.ProviderAdminUserID, credentialID, ) - exec(`INSERT INTO integration_installs(id,org_id,project_id,agent_profile_id,installed_by_user_id, + exec( + `INSERT INTO integration_installs(id,org_id,project_id,agent_profile_id,installed_by_user_id, provider,integration_kind, connection_mode,state,provider_tenant_id,provider_account_ref,provider_identity, credential_secret_id,created_at,updated_at) VALUES($1,$2,$3,$4,$5,'slack','agent_profile','webhook','disabled','T456','A456', '{"bot_user_id":"U456"}',$6,'2026-01-01','2026-01-01')`, - secondIntegrationID, ids.OrgID, ids.ProjectID, profileID, ids.ProviderAdminUserID, credentialID) + secondIntegrationID, + ids.OrgID, + ids.ProjectID, + profileID, + ids.ProviderAdminUserID, + credentialID, + ) agents := []uuid.UUID{uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New()} const archivedIndex, subagentIndex, noTurnIndex = 2, 3, 4 for i, agentID := range agents { @@ -295,14 +352,59 @@ tools: ) } if i == 0 || i == 1 { + openingInput, openingEvent := uuid.New(), uuid.New() + q( + `INSERT INTO agent_inputs(id,project_id,agent_id,state,input_kind,delivery_mode,queued_at) + VALUES($1,$2,$3,'received','content','queued',now())`, + openingInput, + ids.ProjectID, + agentID, + ) + q( + `INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,ordinal,block_kind,text_content,created_at) + VALUES($1,'agent_input',$2,0,'text','Review the conversation',now())`, + agentID, + openingInput, + ) + q( + `INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,idempotency_key,agent_input_id,is_opening_event,created_at) + VALUES($1,$2,$3,2,'agent_input',$4,$5,true,now())`, + openingEvent, + agentID, + turnID, + "agent_input:"+openingInput.String(), + openingInput, + ) + q( + `UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=now(),resolved_at=now() WHERE id=$1`, + openingInput, + openingEvent, + ) + q( + `UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, + turnID, + openingEvent, + ) + q(`UPDATE agents SET next_event_sequence=3 WHERE id=$1`, agentID) failed := uuid.New() - q(`INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,operation_kind,attempt_number, + q( + `INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,operation_kind,attempt_number, agent_config_id,configured_model_revision_id,input_event_sequence, runtime_lock_id,state,created_at) - VALUES($1,$2,$3,$4,'normal',1,$5,$6,1,$7,'started',now())`, - failed, ids.OrgID, ids.ProjectID, agentID, configID, revisionID, uuid.New()) - q(`UPDATE model_call_contexts SET state='failed',recovery_kind='retry',error_kind='provider_error', - retry_at=now(),completed_at=now() WHERE id=$1`, failed) + VALUES($1,$2,$3,$4,'normal',1,$5,$6,2,$7,'started',now())`, + failed, + ids.OrgID, + ids.ProjectID, + agentID, + configID, + revisionID, + uuid.New(), + ) + q( + `UPDATE model_call_contexts SET state='failed',recovery_kind='retry',error_kind='provider_error', + retry_at=now(),completed_at=now() WHERE id=$1`, + failed, + ) if i == 0 && scenario != "continuable_retry" { seedSlackCutoverSuccessfulRetry(t, q, ids, agentID, configID, revisionID, turnID) } @@ -317,7 +419,7 @@ tools: ) q( `INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,idempotency_key,agent_input_id,is_opening_event,created_at) - VALUES($1,$2,$3,2,'agent_input',$4,$5,false,now())`, + VALUES($1,$2,$3,3,'agent_input',$4,$5,false,now())`, stopEvent, agentID, turnID, @@ -329,7 +431,7 @@ tools: stopInput, stopEvent, ) - q(`UPDATE agents SET next_event_sequence=3 WHERE id=$1`, agentID) + q(`UPDATE agents SET next_event_sequence=4 WHERE id=$1`, agentID) q( `UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, turnID, @@ -371,13 +473,19 @@ tools: if scenario == "multiple_targets" { otherIntegration = integrationID } - q(`INSERT INTO integration_targets(id,project_id,agent_id,integration_install_id,target_ref, + q( + `INSERT INTO integration_targets(id,project_id,agent_id,integration_install_id,target_ref, provider_ref,provider_ref_kind,created_at,updated_at) VALUES($1,$2,$3,$4,'other',$5,'dm',now(),now())`, - uuid.New(), ids.ProjectID, agentID, otherIntegration, fmt.Sprintf( + uuid.New(), + ids.ProjectID, + agentID, + otherIntegration, + fmt.Sprintf( "D%d", i+1, - )) + ), + ) } require.NoError(t, tx.Commit()) } @@ -392,7 +500,7 @@ tools: require.NoError(t, db.QueryRowContext(ctx, `SELECT jsonb_build_object( 'agents',(SELECT jsonb_agg(to_jsonb(a)-'current_config_id'-'next_event_sequence' -'integration_target_id'-'interaction_target_id'-'interaction_auto_select'-'updated_at' - -'interaction_handler_key' ORDER BY a.id) FROM agents a), + -'interaction_handler_key'-'root_agent_id' ORDER BY a.id) FROM agents a), 'targets',(SELECT jsonb_agg((to_jsonb(t)-'integration_install_id'-'integration_id' -'launch_key'-'target_ref'-'provider_metadata' -'provider_ref_kind'-'provider_ref'-'scope_kind'-'scope_ref') @@ -401,9 +509,12 @@ tools: 'scope_kind',coalesce(to_jsonb(t)->'scope_kind',to_jsonb(t)->'provider_ref_kind'), 'scope_ref',coalesce(to_jsonb(t)->'scope_ref',to_jsonb(t)->'provider_ref')) ORDER BY t.id) FROM integration_targets t), - 'inputs',(SELECT jsonb_agg(to_jsonb(i) ORDER BY i.id) FROM agent_inputs i + 'inputs',(SELECT jsonb_agg(to_jsonb(i)-'opening_input_ids'-'opening_event_sequence' ORDER + BY i.id) FROM agent_inputs i WHERE i.input_idempotency_key IS DISTINCT FROM 'slack_integration_cutover'), - 'contexts',(SELECT jsonb_agg(to_jsonb(c) ORDER BY c.id) FROM model_call_contexts c), + 'contexts',(SELECT + jsonb_agg(to_jsonb(c)-'turn_id'-'opening_input_ids'-'opening_event_sequence' ORDER BY + c.id) FROM model_call_contexts c), 'outputs',(SELECT jsonb_agg(to_jsonb(o) ORDER BY o.id) FROM model_outputs o), 'blocks',(SELECT jsonb_agg(to_jsonb(b) ORDER BY b.id) FROM content_blocks b), 'secrets',(SELECT jsonb_agg(to_jsonb(s) ORDER BY s.id) FROM secrets s), @@ -414,18 +525,32 @@ tools: case "bad_hash", "bad_source_hash": exec(`ALTER TABLE agent_configs DISABLE TRIGGER agent_configs_immutable`) if scenario == "bad_hash" { - exec(`UPDATE agent_configs SET effective_definition_hash=repeat('0',64) WHERE id=$1`, configID) + exec( + `UPDATE agent_configs SET effective_definition_hash=repeat('0',64) WHERE id=$1`, + configID, + ) } else { exec(`UPDATE agent_configs SET source_hash=repeat('0',64) WHERE id=$1`, configID) } exec(`ALTER TABLE agent_configs ENABLE TRIGGER agent_configs_immutable`) case "preflight_quota": - exec(`INSERT INTO org_resource_limit_overrides(org_id,max_agent_configs_per_project) VALUES($1,1)`, ids.OrgID) + exec( + `INSERT INTO org_resource_limit_overrides(org_id,max_agent_configs_per_project) VALUES($1,1)`, + ids.OrgID, + ) case "live_lease": - exec(`INSERT INTO agent_runtime_locks(agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) - VALUES($1,$2,now(),now(),now()+interval '1 hour')`, agents[0], uuid.New()) + exec( + `INSERT INTO agent_runtime_locks(agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + VALUES($1,$2,now(),now(),now()+interval '1 hour')`, + agents[0], + uuid.New(), + ) case "mislabeled_fixed_agent": - exec(`UPDATE integration_installs SET agent_profile_id=NULL,agent_id=$2 WHERE id=$1`, integrationID, agents[0]) + exec( + `UPDATE integration_installs SET agent_profile_id=NULL,agent_id=$2 WHERE id=$1`, + integrationID, + agents[0], + ) case "fixed_agent": exec( `UPDATE integration_installs SET agent_profile_id=NULL,agent_id=$2,integration_kind='agent' WHERE id=$1`, @@ -452,15 +577,23 @@ tools: VALUES($1,$2,$3,$4::jsonb,$5,now())`, ids.OrgID, projectID, modelID, compiled, hash) } switch scenario { - case "preflight_invalid_policy", "preflight_unmapped_policy", "preflight_address", "preflight_quota", - "custom_collision", "live_lease", - "unsupported_setup", "fixed_agent", "mislabeled_fixed_agent", "multiple_targets": + case "preflight_invalid_policy", + "preflight_unmapped_policy", + "preflight_address", + "preflight_quota", + "custom_collision", + "live_lease", + "unsupported_setup", + "fixed_agent", + "mislabeled_fixed_agent", + "multiple_targets": err := applyProductionPostgresMigrations(ctx, db) want := "requires maintenance" if scenario == "custom_collision" { want = "conflicts with a new integration built-in" } - if scenario == "unsupported_setup" || scenario == "fixed_agent" || scenario == "mislabeled_fixed_agent" { + if scenario == "unsupported_setup" || scenario == "fixed_agent" || + scenario == "mislabeled_fixed_agent" { want = "not a profile-bound Slack webhook setup" } if scenario == "multiple_targets" { @@ -478,8 +611,11 @@ tools: require.ErrorContains(t, err, want) require.Equal(t, int64(47), currentPostgresMigrationVersion(t, ctx, db)) var oldConnections int - require.NoError(t, db.QueryRowContext(ctx, - `SELECT count(*) FROM integration_installs WHERE id=$1`, integrationID).Scan(&oldConnections)) + require.NoError(t, db.QueryRowContext( + ctx, + `SELECT count(*) FROM integration_installs WHERE id=$1`, + integrationID, + ).Scan(&oldConnections)) require.Equal(t, 1, oldConnections) var newTableExists bool require.NoError(t, db.QueryRowContext(ctx, @@ -545,28 +681,42 @@ tools: assertConversationRollback := func() { t.Helper() var conversations int - require.NoError(t, db.QueryRowContext(ctx, - `SELECT count(*) FROM integration_states WHERE kind='agent_conversation'`).Scan(&conversations)) + require.NoError(t, db.QueryRowContext( + ctx, + `SELECT count(*) FROM integration_states WHERE kind='agent_conversation'`, + ).Scan(&conversations)) require.Zero(t, conversations, "migration failure rolls back conversation assignments") var subscriptions, owners int - require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM integration_subscriptions`).Scan(&subscriptions)) + require.NoError( + t, + db.QueryRowContext(ctx, `SELECT count(*) FROM integration_subscriptions`). + Scan(&subscriptions), + ) require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM integration_targets WHERE launch_key IS NOT NULL`).Scan(&owners)) require.Zero(t, subscriptions) require.Zero(t, owners) } if scenario == "invalid_policy" || scenario == "unmapped_policy" || scenario == "invalid_address" || - scenario == "bad_hash" || scenario == "bad_source_hash" || scenario == "config_limit" { + scenario == "bad_hash" || + scenario == "bad_source_hash" || + scenario == "config_limit" { require.NoError(t, applyProductionPostgresMigrationsThrough(t, ctx, db, 48)) if scenario == "config_limit" { - exec(`INSERT INTO org_resource_limit_overrides(org_id,max_agent_configs_per_project) VALUES($1,1)`, ids.OrgID) + exec( + `INSERT INTO org_resource_limit_overrides(org_id,max_agent_configs_per_project) VALUES($1,1)`, + ids.OrgID, + ) } if scenario == "invalid_address" { exec(`UPDATE integration_targets SET scope_ref='C1:invalid' WHERE agent_id=$1`, agents[0]) } expectedConfigs := 1 if scenario == "invalid_policy" { - legacySendPolicy["permission"] = map[string]any{"mode": "always_ask", "parameters": map[string]any{}} + legacySendPolicy["permission"] = map[string]any{ + "mode": "always_ask", + "parameters": map[string]any{}, + } invalid, err := json.Marshal(compiledObject) require.NoError(t, err) exec(`ALTER TABLE agent_configs DISABLE TRIGGER agent_configs_immutable`) @@ -659,7 +809,10 @@ tools: ), ) require.Equal(t, len(agents), active) - require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM agent_configs`).Scan(&configs)) + require.NoError( + t, + db.QueryRowContext(ctx, `SELECT count(*) FROM agent_configs`).Scan(&configs), + ) require.Equal(t, 1, configs) var triggerEnabled string require.NoError( @@ -672,7 +825,9 @@ tools: ), ) require.Equal(t, "O", triggerEnabled) - exec(`DROP TRIGGER reject_cutover_rewrite ON agent_configs; DROP FUNCTION reject_cutover_rewrite()`) + exec( + `DROP TRIGGER reject_cutover_rewrite ON agent_configs; DROP FUNCTION reject_cutover_rewrite()`, + ) } if scenario == "injected_failure" { @@ -694,12 +849,16 @@ tools: require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM agent_configs`).Scan(&count)) require.Equal(t, 1, count) require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM agent_inputs`).Scan(&count)) - require.Equal(t, len(agents)+1, count) - exec(`DROP TRIGGER reject_cutover ON agent_inputs; DROP FUNCTION reject_cutover_config_event()`) + require.Equal(t, len(agents)+3, count) + exec( + `DROP TRIGGER reject_cutover ON agent_inputs; DROP FUNCTION reject_cutover_config_event()`, + ) } if scenario == "pre_internal_ids" { require.NoError(t, applyProductionPostgresMigrationsThrough(t, ctx, db, 48)) - exec(`ALTER TABLE event_webhook_deliveries ADD CONSTRAINT reject_cutover_webhook CHECK (false) NOT VALID`) + exec( + `ALTER TABLE event_webhook_deliveries ADD CONSTRAINT reject_cutover_webhook CHECK (false) NOT VALID`, + ) before := history() err := applyProductionPostgresMigrations(ctx, db) require.ErrorContains(t, err, "reject_cutover_webhook") @@ -707,9 +866,16 @@ tools: assertConversationRollback() require.JSONEq(t, before, history()) var configs, deliveries int - require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM agent_configs`).Scan(&configs)) + require.NoError( + t, + db.QueryRowContext(ctx, `SELECT count(*) FROM agent_configs`).Scan(&configs), + ) require.Equal(t, 1, configs) - require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM event_webhook_deliveries`).Scan(&deliveries)) + require.NoError( + t, + db.QueryRowContext(ctx, `SELECT count(*) FROM event_webhook_deliveries`). + Scan(&deliveries), + ) require.Zero(t, deliveries) exec(`ALTER TABLE event_webhook_deliveries DROP CONSTRAINT reject_cutover_webhook`) } @@ -718,7 +884,11 @@ tools: require.JSONEq(t, before, history()) if scenario != "pre_internal_ids" { var deliveries int - require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM event_webhook_deliveries`).Scan(&deliveries)) + require.NoError( + t, + db.QueryRowContext(ctx, `SELECT count(*) FROM event_webhook_deliveries`). + Scan(&deliveries), + ) require.Zero(t, deliveries, "cutover must not enqueue deliveries without an event webhook") } var retiredSlot sql.NullString @@ -727,7 +897,10 @@ tools: Scan(&retiredSlot)) require.False(t, retiredSlot.Valid) execution := executionstore.New(pool, executionstore.Config{}) - conversations := integrationstore.New(pool, executionstore.IntegrationAccess{}) + conversations := integrationstore.New( + pool, + integrationstore.Config{Access: executionstore.IntegrationAccess{}}, + ) for i, agentID := range agents { snapshot, err := execution.CaptureAgentConfigForModelContext(ctx, ids.ProjectID, agentID) require.NoError(t, err) @@ -742,6 +915,9 @@ tools: if i == noTurnIndex { expectedSequence = 1 } + if i == 0 || i == 1 { + expectedSequence++ + } require.Equal(t, expectedSequence, snapshot.InputEventSequence) contract, err := agentconfig.RuntimeContractFromCompiled( snapshot.AgentConfig.CompiledDefinition, @@ -753,8 +929,11 @@ tools: require.NoError(t, json.Unmarshal(snapshot.AgentConfig.CompiledDefinition, &raw)) if scenario == "pre_internal_ids" { var sequence int64 - require.NoError(t, db.QueryRowContext(ctx, - `SELECT event_sequence FROM event_webhook_deliveries WHERE agent_id=$1`, agentID).Scan(&sequence)) + require.NoError(t, db.QueryRowContext( + ctx, + `SELECT event_sequence FROM event_webhook_deliveries WHERE agent_id=$1`, + agentID, + ).Scan(&sequence)) require.Equal(t, snapshot.InputEventSequence, sequence) require.NotNil(t, raw.EventWebhook) require.Equal(t, ids.ProviderSecretID, raw.EventWebhook.SigningSecretID) @@ -778,26 +957,52 @@ tools: ) require.NoError(t, err) require.True(t, found) - wantConversation := integrationstore.ConversationAddress{Kind: "thread", Ref: fmt.Sprintf("C123:111.%d", i+1)} + wantConversation := integrationstore.ConversationAddress{ + Kind: "thread", + Ref: fmt.Sprintf("C123:111.%d", i+1), + } if scenario == "channel" { - wantConversation = integrationstore.ConversationAddress{Kind: "channel", Ref: fmt.Sprintf("C%d", i+1)} + wantConversation = integrationstore.ConversationAddress{ + Kind: "channel", + Ref: fmt.Sprintf("C%d", i+1), + } } if scenario == "dm" { - wantConversation = integrationstore.ConversationAddress{Kind: "dm", Ref: fmt.Sprintf("D%d", i+1)} + wantConversation = integrationstore.ConversationAddress{ + Kind: "dm", + Ref: fmt.Sprintf("D%d", i+1), + } } require.Equal(t, wantConversation, conversation) - assertSlackCutoverConversationState(t, db, ids.ProjectID, agentID, integrationID, wantConversation) + assertSlackCutoverConversationState( + t, + db, + ids.ProjectID, + agentID, + integrationID, + wantConversation, + ) if scenario == "normal" || scenario == "dm" { tx, err := pool.Begin(ctx) require.NoError(t, err) - candidates, routeErr := conversations.IntegrationRoutingCandidatesTx(ctx, tx, - ids.ProjectID, integrationID, wantConversation, []integrationstore.ConversationAddress{wantConversation}) + candidates, routeErr := conversations.IntegrationRoutingCandidatesTx( + ctx, + tx, + ids.ProjectID, + integrationID, + wantConversation, + []integrationstore.ConversationAddress{wantConversation}, + ) require.NoError(t, tx.Rollback(ctx)) require.NoError(t, routeErr) require.NotNil(t, candidates.Launcher) require.Len(t, candidates.LaunchOwners, 1) require.Equal(t, agentID, candidates.LaunchOwners[0].AgentID) - require.Equal(t, integrationdefinition.ProfileLaunchKey, candidates.LaunchOwners[0].LaunchKey) + require.Equal( + t, + integrationdefinition.ProfileLaunchKey, + candidates.LaunchOwners[0].LaunchKey, + ) if i == archivedIndex { require.Empty(t, candidates.Subscriptions) } else { @@ -807,17 +1012,28 @@ tools: scope, err := integrationdefinition.ParseConversation(integrationdefinition.ProviderSlack, wantConversation.Kind, wantConversation.Ref) require.NoError(t, err) - event := integrationdefinition.Event{Scope: scope, Kind: integrationdefinition.EventMessage, Mentioned: true} + event := integrationdefinition.Event{ + Scope: scope, + Kind: integrationdefinition.EventMessage, + Mentioned: true, + } definition, ok := integrationdefinition.Lookup(candidates.Launcher.IntegrationKind) require.True(t, ok) require.True(t, definition.MatchesLaunch(candidates.Launcher.Settings, event)) launcher := integration.NewChatIntegrationLauncher(conversations, execution, nil) launches, err := launcher.Decide(ctx, integration.IntegrationLaunchContext{ Integration: *candidates.Launcher, Address: wantConversation, Candidates: candidates, - Event: integration.IntegrationEvent{Event: event, SemanticKey: "post-migration-mention"}, + Event: integration.IntegrationEvent{ + Event: event, + SemanticKey: "post-migration-mention", + }, }) require.NoError(t, err) - require.Empty(t, launches, "mentioning an existing conversation must not launch a replacement agent") + require.Empty( + t, + launches, + "mentioning an existing conversation must not launch a replacement agent", + ) } wantIntegrations := 1 if scenario == "multiple_integrations" { @@ -830,9 +1046,19 @@ tools: ) require.NoError(t, err) require.True(t, found) - wantSecond := integrationstore.ConversationAddress{Kind: "dm", Ref: fmt.Sprintf("D%d", i+1)} + wantSecond := integrationstore.ConversationAddress{ + Kind: "dm", + Ref: fmt.Sprintf("D%d", i+1), + } require.Equal(t, wantSecond, second) - assertSlackCutoverConversationState(t, db, ids.ProjectID, agentID, secondIntegrationID, wantSecond) + assertSlackCutoverConversationState( + t, + db, + ids.ProjectID, + agentID, + secondIntegrationID, + wantSecond, + ) } require.Len(t, contract.IntegrationTools, wantIntegrations) require.Contains(t, raw.Tools, "set_interaction_handler") @@ -847,7 +1073,12 @@ tools: WHERE event.agent_id=$1 AND event.sequence=1`, agentID).Scan(&oldID), ) require.Equal(t, configID.String(), oldID) - historical, err := execution.CaptureAgentConfigForEventWatermark(ctx, ids.ProjectID, agentID, 1) + historical, err := execution.CaptureAgentConfigForEventWatermark( + ctx, + ids.ProjectID, + agentID, + 1, + ) require.NoError(t, err) require.Equal(t, configID, historical.AgentConfig.ID) } @@ -856,7 +1087,7 @@ tools: t, db.QueryRowContext( ctx, - `SELECT EXISTS(SELECT 1 FROM agent_next_model_work($1,$2))`, + `SELECT EXISTS(SELECT 1 FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND h.agent_id=$2 AND h.logical_ready_at IS NOT NULL)`, ids.ProjectID, agentID, ). @@ -864,7 +1095,7 @@ tools: &runnable, ), ) - require.False(t, runnable) + require.Equal(t, scenario == "continuable_retry" && i == 0, runnable) } rows, err := db.QueryContext( ctx, @@ -897,7 +1128,11 @@ tools: } require.Equal(t, wantOwners, selections, "existing threads and DMs keep their launch owner") var subscriptions, pointers int - require.NoError(t, db.QueryRowContext(ctx, `SELECT count(*) FROM integration_subscriptions`).Scan(&subscriptions)) + require.NoError( + t, + db.QueryRowContext(ctx, `SELECT count(*) FROM integration_subscriptions`). + Scan(&subscriptions), + ) wantSubscriptions := len(agents) - 1 if scenario == "multiple_integrations" { wantSubscriptions *= 2 @@ -905,7 +1140,12 @@ tools: if scenario == "channel" { wantSubscriptions = 0 } - require.Equal(t, wantSubscriptions, subscriptions, "active agents retain their thread and DM subscriptions") + require.Equal( + t, + wantSubscriptions, + subscriptions, + "active agents retain their thread and DM subscriptions", + ) require.NoError( t, db.QueryRowContext( @@ -983,17 +1223,23 @@ tools: require.False(t, tool.Enabled) } } - launched, err := agentconfig.DeriveWithIntegrationCapabilities(historical, agentconfig.IntegrationCapabilitiesSource{ - Tools: map[string]agentconfig.AgentConfigToolSource{ - "int__slack__post_message": {}, + launched, err := agentconfig.DeriveWithIntegrationCapabilities( + historical, + agentconfig.IntegrationCapabilitiesSource{ + Tools: map[string]agentconfig.AgentConfigToolSource{ + "int__slack__post_message": {}, + }, + }, + agentconfig.CompileOptions{ + ResolveIntegrationName: func(name string) (agentconfig.IntegrationResolution, error) { + require.Equal(t, "slack", name) + return agentconfig.IntegrationResolution{ + IntegrationID: integrationID, + IntegrationKind: integrationdefinition.SlackThread, + }, nil + }, }, - }, agentconfig.CompileOptions{ResolveIntegrationName: func(name string) (agentconfig.IntegrationResolution, error) { - require.Equal(t, "slack", name) - return agentconfig.IntegrationResolution{ - IntegrationID: integrationID, - IntegrationKind: integrationdefinition.SlackThread, - }, nil - }}) + ) require.NoError(t, err) if sendingEnabled { require.True(t, launched.Tools["int__slack__post_message"].Enabled) @@ -1005,7 +1251,11 @@ tools: require.Equal(t, source != nil, migratedSource.Valid) if source != nil { require.Equal(t, sourceFormat, migratedFormat.String) - require.Equal(t, fmt.Sprintf("%x", sha256.Sum256([]byte(migratedSource.String))), migratedSourceHash.String) + require.Equal( + t, + fmt.Sprintf("%x", sha256.Sum256([]byte(migratedSource.String))), + migratedSourceHash.String, + ) var object map[string]any if sourceFormat == "json" { require.NoError(t, json.Unmarshal([]byte(migratedSource.String), &object)) @@ -1069,15 +1319,21 @@ tools: require.NoError(t, err) require.True(t, created) if scenario == "pre_internal_ids" && i == 0 { - claim, found, err := execution.ClaimNextAgentWork(ctx, executionstore.ClaimNextAgentWorkInput{ - WorkerProcessID: uuid.New(), LeaseDuration: time.Minute, - }) + claim, found, err := execution.ClaimNextAgentWork( + ctx, + executionstore.ClaimNextAgentWorkInput{ + WorkerProcessID: uuid.New(), LeaseDuration: time.Minute, + }, + ) require.NoError(t, err) require.True(t, found) require.Equal(t, executionstore.AgentWorkModel, claim.Kind) var deliveries int - require.NoError(t, db.QueryRowContext(ctx, - `SELECT count(*) FROM event_webhook_deliveries WHERE agent_id=$1`, agentID).Scan(&deliveries)) + require.NoError(t, db.QueryRowContext( + ctx, + `SELECT count(*) FROM event_webhook_deliveries WHERE agent_id=$1`, + agentID, + ).Scan(&deliveries)) require.Equal(t, 2, deliveries) } snapshot, err := execution.CaptureAgentConfigForModelContext(ctx, ids.ProjectID, agentID) @@ -1085,7 +1341,14 @@ tools: require.NotEqual(t, configID, snapshot.AgentConfig.ID) } if scenario == "normal" { - assertSlackCutoverInputGuards(t, db, ids.ProjectID, agents[noTurnIndex], agents[0], integrationID) + assertSlackCutoverInputGuards( + t, + db, + ids.ProjectID, + agents[noTurnIndex], + agents[0], + integrationID, + ) } }) } @@ -1099,41 +1362,63 @@ func assertSlackCutoverInputGuards( t.Helper() ctx := t.Context() var targetID, inputID uuid.UUID - require.NoError(t, db.QueryRowContext(ctx, + require.NoError(t, db.QueryRowContext( + ctx, `SELECT id FROM integration_targets WHERE project_id=$1 AND agent_id=$2 AND integration_id=$3 AND deleted_at IS NULL`, - projectID, agentID, integrationID).Scan(&targetID)) - require.NoError(t, db.QueryRowContext(ctx, + projectID, + agentID, + integrationID, + ).Scan(&targetID)) + require.NoError(t, db.QueryRowContext( + ctx, `INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,integration_target_id,queued_at) VALUES($1,$2,'received','content','queued',$3,now()) RETURNING id`, - projectID, agentID, targetID).Scan(&inputID)) + projectID, + agentID, + targetID, + ).Scan(&inputID)) _, err := db.ExecContext(ctx, `UPDATE agent_inputs SET delivery_mode='steering' WHERE id=$1`, inputID) require.NoError(t, err, "received content remains editable after cutover") _, err = db.ExecContext(ctx, `UPDATE agent_inputs SET integration_target_id=NULL WHERE id=$1`, inputID) require.ErrorContains(t, err, "agent_input intent and identity are immutable") - _, err = db.ExecContext(ctx, `UPDATE agent_inputs SET state='canceled',canceled_at=now() WHERE id=$1`, inputID) + _, err = db.ExecContext( + ctx, + `UPDATE agent_inputs SET state='canceled',canceled_at=now() WHERE id=$1`, + inputID, + ) require.NoError(t, err) _, err = db.ExecContext(ctx, `UPDATE agent_inputs SET integration_target_id=NULL WHERE id=$1`, inputID) require.ErrorContains(t, err, "terminal agent_inputs are immutable") _, err = db.ExecContext(ctx, `DELETE FROM agent_inputs WHERE id=$1`, inputID) require.ErrorContains(t, err, "agent_inputs are immutable") - _, err = db.ExecContext(ctx, + _, err = db.ExecContext( + ctx, `INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,integration_target_id,queued_at) - VALUES($1,$2,'received','content','queued',$3,now())`, projectID, otherAgentID, targetID) + VALUES($1,$2,'received','content','queued',$3,now())`, + projectID, + otherAgentID, + targetID, + ) require.ErrorContains(t, err, "agent_inputs_project_id_agent_id_integration_target_id_fkey") _, err = db.ExecContext(ctx, `UPDATE agents SET interaction_target_id=$2,interaction_handler_key='int__slack__default' WHERE id=$1`, otherAgentID, targetID) require.ErrorContains(t, err, "agents_project_id_id_interaction_target_id_fkey") - _, err = db.ExecContext(ctx, + _, err = db.ExecContext( + ctx, `INSERT INTO integration_targets(project_id,agent_id,integration_id,scope_kind,scope_ref,created_at,updated_at) SELECT project_id,agent_id,integration_id,scope_kind,scope_ref,now(),now() FROM integration_targets WHERE id=$1`, - targetID) + targetID, + ) require.ErrorContains(t, err, "integration_targets_active_agent_address_idx") } func assertSlackCutoverConversationState( - t *testing.T, db *sql.DB, projectID, agentID, integrationID uuid.UUID, want integrationstore.ConversationAddress, + t *testing.T, + db *sql.DB, + projectID, agentID, integrationID uuid.UUID, + want integrationstore.ConversationAddress, ) { t.Helper() var data []byte @@ -1162,7 +1447,7 @@ func seedSlackCutoverSuccessfulRetry( exec(`INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,operation_kind,attempt_number, agent_config_id,configured_model_revision_id,input_event_sequence, runtime_lock_id,state,created_at) - VALUES($1,$2,$3,$4,'normal',2,$5,$6,1,$7,'started',now())`, + VALUES($1,$2,$3,$4,'normal',2,$5,$6,2,$7,'started',now())`, succeeded, ids.OrgID, ids.ProjectID, agentID, configID, revisionID, uuid.New()) exec( `INSERT INTO model_outputs(id,agent_id,model_call_context_id,stop_reason,created_at) VALUES($1,$2,$3,'end_turn',now())`, @@ -1170,21 +1455,25 @@ func seedSlackCutoverSuccessfulRetry( agentID, succeeded, ) - exec(`INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,text_content,created_at) - VALUES($1,'model_output',$2,1,'text','Historical Slack reply',now())`, agentID, outputID) + exec( + `INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,text_content,created_at) + VALUES($1,'model_output',$2,1,'text','Historical Slack reply',now())`, + agentID, + outputID, + ) exec( `UPDATE model_call_contexts SET state='succeeded',api_format='openai',api_variant='responses',completed_at=now() WHERE id=$1`, succeeded, ) exec( `INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,idempotency_key,model_output_id,is_opening_event,created_at) - VALUES($1,$2,$3,2,'model_output',$4,$5,false,now())`, + VALUES($1,$2,$3,3,'model_output',$4,$5,false,now())`, eventID, agentID, turnID, "model_output:"+outputID.String(), outputID, ) - exec(`UPDATE agents SET next_event_sequence=3 WHERE id=$1`, agentID) + exec(`UPDATE agents SET next_event_sequence=4 WHERE id=$1`, agentID) exec(`UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, turnID, eventID) } diff --git a/internal/dbmigrate/slack_integration_pending_work_integration_test.go b/internal/dbmigrate/slack_integration_pending_work_integration_test.go index c64438a23..99ae9cdb9 100644 --- a/internal/dbmigrate/slack_integration_pending_work_integration_test.go +++ b/internal/dbmigrate/slack_integration_pending_work_integration_test.go @@ -46,24 +46,40 @@ func TestSlackIntegrationCutoverPendingInteractionsRecover(t *testing.T) { f.migrate(t) ctx := t.Context() execution := f.store.Execution() - interaction, found, err := execution.GetAgentInteraction(ctx, f.ids.ProjectID, f.agentID, f.interactionID) + interaction, found, err := execution.GetAgentInteraction( + ctx, + f.ids.ProjectID, + f.agentID, + f.interactionID, + ) require.NoError(t, err) require.True(t, found) require.Equal(t, executionstore.AgentInteractionStateOpen, interaction.State) require.Equal(t, f.turnID, interaction.TurnID) require.JSONEq(t, string(f.request), string(interaction.Request)) - require.Empty(t, interaction.Destination, "old interactions keep their original dashboard resolution path") + require.Empty( + t, + interaction.Destination, + "old interactions keep their original dashboard resolution path", + ) _, found, err = execution.ClaimNextAgentWork(ctx, pendingWorkClaimInput()) require.ErrorIs(t, err, storeerr.ErrNoClaimableAgentWakeup) require.False(t, found, "an unanswered interaction must stay parked after config cutover") resolution := executionstore.ResolveAgentInteractionInput{ ProjectID: f.ids.ProjectID, AgentID: f.agentID, ID: f.interactionID, Actor: f.actor, - Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{OptionIndices: []int{0}}}}, + Resolution: interactionform.Resolution{ + Answers: []interactionform.Answer{{OptionIndices: []int{0}}}, + }, } if scenario.steer { input := f.steer(t) - interaction, found, err = execution.GetAgentInteraction(ctx, f.ids.ProjectID, f.agentID, f.interactionID) + interaction, found, err = execution.GetAgentInteraction( + ctx, + f.ids.ProjectID, + f.agentID, + f.interactionID, + ) require.NoError(t, err) require.True(t, found) require.Equal(t, executionstore.AgentInteractionStateCanceled, interaction.State) @@ -124,7 +140,12 @@ func TestSlackIntegrationCutoverExpiredRuntimeRecovers(t *testing.T) { require.True(t, found) require.Equal(t, executionstore.ModelCallContextStarted, model.State) require.Equal(t, f.runtimeID, model.RuntimeLockID) - require.Equal(t, f.configID, model.AgentConfigID, "cutover must not rewrite an in-flight context's config identity") + require.Equal( + t, + f.configID, + model.AgentConfigID, + "cutover must not rewrite an in-flight context's config identity", + ) reaped, err := execution.ReapExpiredAgentRuntimeLocks(ctx, 10) require.NoError(t, err) require.EqualValues(t, 1, reaped) @@ -175,12 +196,20 @@ func (f *slackPendingWorkFixture) migrate(t *testing.T) { } var turnID, configID uuid.UUID var opening bool - require.NoError(t, f.db.QueryRowContext(ctx, `SELECT event.turn_id,event.is_opening_event,input.agent_config_id + require.NoError( + t, + f.db.QueryRowContext(ctx, `SELECT event.turn_id,event.is_opening_event,input.agent_config_id FROM agent_inputs input JOIN agent_events event ON event.agent_id=input.agent_id AND event.id=input.admitted_event_id WHERE input.agent_id=$1 AND input.input_idempotency_key='slack_integration_cutover'`, f.agentID). - Scan(&turnID, &opening, &configID)) - require.Equal(t, f.turnID, turnID, "successor config must be admitted into the pending work's original turn") + Scan(&turnID, &opening, &configID), + ) + require.Equal( + t, + f.turnID, + turnID, + "successor config must be admitted into the pending work's original turn", + ) require.False(t, opening) require.NotEqual(t, f.configID, configID) var owner, pinnedTarget uuid.UUID @@ -189,7 +218,11 @@ func (f *slackPendingWorkFixture) migrate(t *testing.T) { agent.interaction_handler_key FROM integration_targets target JOIN agents agent ON agent.id=target.agent_id WHERE target.id=$1 AND target.integration_id=$2 AND target.scope_kind=$3 AND target.scope_ref=$4 - AND target.launch_key='default' AND target.deleted_at IS NULL`, f.targetID, f.integrationID, f.scope, f.ref). + AND target.launch_key='default' AND target.deleted_at IS NULL`, + f.targetID, + f.integrationID, + f.scope, + f.ref). Scan(&owner, &pinnedTarget, &handler)) require.Equal(t, f.agentID, owner) require.Equal(t, f.targetID, pinnedTarget) @@ -213,12 +246,15 @@ func (f *slackPendingWorkFixture) migrate(t *testing.T) { `SELECT count(*) FROM agent_turns WHERE agent_id=$1`, f.agentID).Scan(&turns)) require.Equal(t, 1, turns) var wakeupCovered bool - require.NoError(t, f.db.QueryRowContext(ctx, `SELECT - (agent_next_wakeup_ready_at(project_id, - id) IS NULL AND NOT EXISTS (SELECT 1 FROM agent_wakeups WHERE agent_id=agents.id)) - OR EXISTS (SELECT 1 FROM agent_wakeups WHERE agent_id=agents.id - AND ready_at=agent_next_wakeup_ready_at(project_id,agents.id)) - FROM agents WHERE id=$1`, f.agentID).Scan(&wakeupCovered)) + require.NoError( + t, + f.db.QueryRowContext(ctx, + `SELECT CASE WHEN EXISTS(SELECT 1 FROM agent_runtime_locks WHERE agent_id=head.agent_id) + THEN wake.agent_id IS NULL WHEN head.logical_ready_at IS NULL THEN wake.agent_id IS NULL + ELSE wake.agent_id IS NOT NULL AND wake.ready_at <= head.logical_ready_at END + FROM agent_execution_state head LEFT JOIN agent_wakeups wake ON wake.agent_id=head.agent_id + WHERE head.agent_id=$1`, f.agentID).Scan(&wakeupCovered), + ) require.True(t, wakeupCovered, "cutover must reconcile the scheduler index with the actual pending work") } @@ -230,10 +266,19 @@ func newSlackPendingWorkFixture(t *testing.T, kind, scope string) slackPendingWo t.Cleanup(func() { _ = db.Close() }) require.NoError(t, applyProductionPostgresMigrationsThrough(t, ctx, db, 47)) f := slackPendingWorkFixture{ - db: db, - store: storage.NewStore(pool, storage.WithModelCallRetryBackoff(func(int, string) time.Duration { return 0 })), - ids: storagefixture.ProjectIDs{OrgID: uuid.New(), ProjectID: uuid.New(), ProviderAdminUserID: uuid.New(), - ProviderSecretID: uuid.New(), ProviderSecretVersionID: uuid.New(), ProviderConfigID: uuid.New()}, + db: db, + store: storage.NewStore( + pool, + storage.WithModelCallRetryBackoff(func(int, string) time.Duration { return 0 }), + ), + ids: storagefixture.ProjectIDs{ + OrgID: uuid.New(), + ProjectID: uuid.New(), + ProviderAdminUserID: uuid.New(), + ProviderSecretID: uuid.New(), + ProviderSecretVersionID: uuid.New(), + ProviderConfigID: uuid.New(), + }, agentID: uuid.New(), turnID: uuid.New(), configID: uuid.New(), contextID: uuid.New(), runtimeID: uuid.New(), integrationID: uuid.New(), targetID: uuid.New(), toolID: uuid.New(), interactionID: uuid.New(), scope: scope, ref: "C123:111.222", @@ -280,7 +325,8 @@ func newSlackPendingWorkFixture(t *testing.T, kind, scope string) slackPendingWo compiled := fmt.Sprintf(`{"instruction":"Review","model":{"configured_model_id":%q},"tools":{ "send_integration_message":{"enabled":true,"permission":{"mode":"always_allow","parameters":{}}}, "ask_question":{"enabled":true,"permission":{"mode":"always_allow","parameters":{}}}, - "read_process":{"enabled":true,"permission":{"mode":"always_ask","parameters":{}}}}}`, modelID.String()) + "read_process":{"enabled":true,"permission":{"mode":"always_ask","parameters":{}}}}}`, + modelID.String()) var compiledObject map[string]any require.NoError(t, json.Unmarshal([]byte(compiled), &compiledObject)) canonical, err := json.Marshal(compiledObject) @@ -295,57 +341,105 @@ func newSlackPendingWorkFixture(t *testing.T, kind, scope string) slackPendingWo VALUES($1,$2,'reviewer',$3,now(),now())) INSERT INTO agent_profile_versions(id,project_id,profile_id,generation,agent_config_id,created_at) VALUES($3,$2,$1,1,$4,now())`, profileID, f.ids.ProjectID, versionID, f.configID) - exec(`INSERT INTO integration_installs(id,org_id,project_id,agent_profile_id,installed_by_user_id,provider, + exec( + `INSERT INTO integration_installs(id,org_id,project_id,agent_profile_id,installed_by_user_id,provider, integration_kind, connection_mode,state,provider_tenant_id,provider_account_ref,provider_identity,credential_secret_id, created_at,updated_at) VALUES($1,$2,$3,$4,$5,'slack','agent_profile','webhook','active','T123','A123', '{"bot_user_id":"U123"}',$6,now(),now())`, - f.integrationID, f.ids.OrgID, f.ids.ProjectID, profileID, f.ids.ProviderAdminUserID, credentialID) + f.integrationID, + f.ids.OrgID, + f.ids.ProjectID, + profileID, + f.ids.ProviderAdminUserID, + credentialID, + ) exec(`INSERT INTO agents(id,org_id,project_id,state,name,agent_profile_id,current_config_id, next_event_sequence,created_at,updated_at) VALUES($1,$2,$3,'active','reviewer',$4,$5,3,now(),now())`, f.agentID, f.ids.OrgID, f.ids.ProjectID, profileID, f.configID) - exec(`INSERT INTO agent_inputs(id,project_id,agent_id,state,input_kind,delivery_mode,agent_config_id,queued_at) + exec( + `INSERT INTO agent_inputs(id,project_id,agent_id,state,input_kind,delivery_mode,agent_config_id,queued_at) VALUES($1,$2,$3,'received','config_change','immediate',$4,now())`, - inputID, f.ids.ProjectID, f.agentID, f.configID) + inputID, + f.ids.ProjectID, + f.agentID, + f.configID, + ) exec(`INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,idempotency_key,agent_input_id, is_opening_event,created_at) VALUES($1,$2,$3,1,'agent_input',$4,$5,true,now())`, eventID, f.agentID, f.turnID, "agent_input:"+inputID.String(), inputID) exec(`INSERT INTO agent_turns(id,agent_id,turn_sequence,latest_event_id,latest_semantic_event_id) VALUES($1,$2,1,$3,$3)`, f.turnID, f.agentID, eventID) - exec(`UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=now(),resolved_at=now() WHERE id=$1`, - inputID, eventID) + exec( + `UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=now(),resolved_at=now() WHERE id=$1`, + inputID, + eventID, + ) contentID, contentEventID := uuid.New(), uuid.New() exec(`INSERT INTO agent_inputs(id,project_id,agent_id,state,input_kind,delivery_mode,queued_at) VALUES($1,$2,$3,'received','content','queued',now())`, contentID, f.ids.ProjectID, f.agentID) - exec(`INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,ordinal,block_kind,text_content,created_at) - VALUES($1,'agent_input',$2,0,'text','Review the change',now())`, f.agentID, contentID) + exec( + `INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,ordinal,block_kind,text_content,created_at) + VALUES($1,'agent_input',$2,0,'text','Review the change',now())`, + f.agentID, + contentID, + ) exec(`INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,idempotency_key,agent_input_id, is_opening_event,created_at) VALUES($1,$2,$3,2,'agent_input',$4,$5,true,now())`, contentEventID, f.agentID, f.turnID, "agent_input:"+contentID.String(), contentID) - exec(`UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=now(),resolved_at=now() WHERE id=$1`, - contentID, contentEventID) - exec(`UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, f.turnID, contentEventID) + exec( + `UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=now(),resolved_at=now() WHERE id=$1`, + contentID, + contentEventID, + ) + exec( + `UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, + f.turnID, + contentEventID, + ) if scope != "" { - exec(`INSERT INTO integration_targets(id,project_id,agent_id,integration_install_id,target_ref,provider_ref, + exec( + `INSERT INTO integration_targets(id,project_id,agent_id,integration_install_id,target_ref,provider_ref, provider_ref_kind,created_at,updated_at) VALUES($1,$2,$3,$4,'slack',$5,$6,now(),now())`, - f.targetID, f.ids.ProjectID, f.agentID, f.integrationID, f.ref, scope) + f.targetID, + f.ids.ProjectID, + f.agentID, + f.integrationID, + f.ref, + scope, + ) exec(`UPDATE agents SET integration_target_id=$2 WHERE id=$1`, f.agentID, f.targetID) - exec(`INSERT INTO agent_wakeups(agent_id,ready_at,updated_at,metadata) VALUES($1,now(),now(),'{}')`, f.agentID) + exec( + `INSERT INTO agent_wakeups(agent_id,ready_at,updated_at,metadata) VALUES($1,now(),now(),'{}')`, + f.agentID, + ) } - exec(`INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,operation_kind,attempt_number,agent_config_id, + exec( + `INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,operation_kind,attempt_number,agent_config_id, configured_model_revision_id,input_event_sequence,runtime_lock_id,state,created_at) VALUES($1,$2,$3,$4,'normal',1,$5,$6,2,$7,'started',now())`, - f.contextID, f.ids.OrgID, f.ids.ProjectID, f.agentID, f.configID, revisionID, f.runtimeID) + f.contextID, + f.ids.OrgID, + f.ids.ProjectID, + f.agentID, + f.configID, + revisionID, + f.runtimeID, + ) if kind == "" { - exec(`INSERT INTO agent_runtime_locks(id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + exec( + `INSERT INTO agent_runtime_locks(id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) VALUES($1,$2,$3,now()-interval '2 minutes',now()-interval '2 minutes',now()-interval '1 minute')`, - f.runtimeID, f.agentID, uuid.New()) + f.runtimeID, + f.agentID, + uuid.New(), + ) } else { f.seedInteraction(t, exec, kind) } @@ -360,11 +454,22 @@ func (f *slackPendingWorkFixture) seedInteraction(t *testing.T, exec func(string outputID, eventID := uuid.New(), uuid.New() exec(`INSERT INTO model_outputs(id,agent_id,model_call_context_id,stop_reason,created_at) VALUES($1,$2,$3,'tool_use',now())`, outputID, f.agentID, f.contextID) - exec(`INSERT INTO tool_calls(id,agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) + exec( + `INSERT INTO tool_calls(id,agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) VALUES($1,$2,$3,'pending-review',$4,$5,'built_in','awaiting_authorization',now())`, - f.toolID, f.agentID, outputID, toolName, toolInput) - exec(`INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,tool_call_id,created_at) - VALUES($1,'model_output',$2,0,'tool_call',$3,now())`, f.agentID, outputID, f.toolID) + f.toolID, + f.agentID, + outputID, + toolName, + toolInput, + ) + exec( + `INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,tool_call_id,created_at) + VALUES($1,'model_output',$2,0,'tool_call',$3,now())`, + f.agentID, + outputID, + f.toolID, + ) exec(`UPDATE model_call_contexts SET state='succeeded',api_format='openai',api_variant='responses', completed_at=now() WHERE id=$1`, f.contextID) @@ -373,7 +478,11 @@ func (f *slackPendingWorkFixture) seedInteraction(t *testing.T, exec func(string VALUES($1,$2,$3,3,'model_output',$4,$5,false,now())`, eventID, f.agentID, f.turnID, "model_output:"+outputID.String(), outputID) exec(`UPDATE agents SET next_event_sequence=4 WHERE id=$1`, f.agentID) - exec(`UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, f.turnID, eventID) + exec( + `UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1`, + f.turnID, + eventID, + ) exec(`INSERT INTO agent_interactions(id,agent_id,tool_call_id,interaction_kind,state,request,created_at) VALUES($1,$2,$3,$4,'open',$5,now())`, f.interactionID, f.agentID, f.toolID, kind, f.request) if kind == "permission" { @@ -402,7 +511,10 @@ func pendingWorkRequest(t *testing.T, kind string) (json.RawMessage, string, jso require.NoError(t, err) form, err := toolpermission.NewAllowDenyForm("Allow reading the process?", nil) require.NoError(t, err) - mode, found := toolpermission.FindMode(toolpermission.CommonModeDescriptors(), toolpermission.ModeAlwaysAsk) + mode, found := toolpermission.FindMode( + toolpermission.CommonModeDescriptors(), + toolpermission.ModeAlwaysAsk, + ) require.True(t, found) request, err := toolpermission.NewRequest( mode, toolpermission.DefaultSelection(toolpermission.ModeAlwaysAsk), authorization, form, @@ -433,18 +545,25 @@ func (f *slackPendingWorkFixture) steer(t *testing.T) executionstore.AgentInputR }) require.NoError(t, err) require.True(t, created) - receipt, found, err := integrations.ClaimIntegrationInbox(ctx, integrationstore.ClaimIntegrationInboxInput{ - ProjectID: f.ids.ProjectID, IntegrationID: f.integrationID, LeaseDuration: time.Minute, - }) + receipt, found, err := integrations.ClaimIntegrationInbox( + ctx, + integrationstore.ClaimIntegrationInboxInput{ + ProjectID: f.ids.ProjectID, IntegrationID: f.integrationID, LeaseDuration: time.Minute, + }, + ) require.NoError(t, err) require.True(t, found) router := integration.NewIntegrationRouter(f.store.Execution(), integrations) plan, created, err := router.Freeze(ctx, receipt.Lease(), &integration.IntegrationEvent{ Event: integrationdefinition.Event{ - Kind: integrationdefinition.EventMessage, Mentioned: false, Scope: integrationdefinition.Scope{Slack: &scope}, + Kind: integrationdefinition.EventMessage, + Mentioned: false, + Scope: integrationdefinition.Scope{Slack: &scope}, }, Actor: actor, SemanticKey: "cutover-steering", DeliveryMode: executionstore.DeliveryModeSteering, - ContentBlocks: json.RawMessage(`[{"type":"text","text":"Use the new instructions instead"}]`), + ContentBlocks: json.RawMessage( + `[{"type":"text","text":"Use the new instructions instead"}]`, + ), CancelOpenInteractions: true, }, nil) require.NoError(t, err) @@ -477,21 +596,22 @@ func (f *slackPendingWorkFixture) claimModel( var sequence int64 require.NoError(t, f.db.QueryRowContext(ctx, `SELECT max(sequence) FROM agent_events WHERE agent_id=$1`, f.agentID).Scan(&sequence)) - snapshot, err := f.store.Execution().CaptureAgentConfigForModelContext(ctx, f.ids.ProjectID, f.agentID) - require.NoError(t, err) - claim, err := f.store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: f.ids.ProjectID, AgentID: f.agentID, RuntimeLockID: work.RuntimeLock.ID, - AgentConfigID: snapshot.AgentConfig.ID, InputEventSequence: sequence, - OpeningInputIDs: work.Model.InputIDs, - SourceModelCallContextID: work.Model.SourceModelCallContextID, - SourceModelOutputID: work.Model.SourceModelOutputID, - }) + + prepared1, err := f.store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: f.ids.ProjectID, AgentID: f.agentID, RuntimeLockID: work.RuntimeLock.ID, + OpeningInputIDs: work.Model.InputIDs, + SourceModelCallContextID: work.Model.SourceModelCallContextID, + SourceModelOutputID: work.Model.SourceModelOutputID, + }) + claim := prepared1.Claim + require.NoError(t, err) require.True(t, claim.Created) require.True(t, claim.Claimed) var turnID uuid.UUID require.NoError(t, f.db.QueryRowContext(ctx, - `SELECT turn_id FROM model_call_context_turns WHERE model_call_context_id=$1`, claim.Context.ID).Scan(&turnID)) + `SELECT turn_id FROM model_call_contexts WHERE id=$1`, claim.Context.ID).Scan(&turnID)) require.Equal(t, work.Model.TurnID, turnID) return claim.Context } @@ -535,7 +655,9 @@ func (f *slackPendingWorkFixture) captureNewPrompts(t *testing.T, work execution question, err := f.store.Execution().ExecuteToolCall(ctx, executionstore.ExecuteToolCallInput{ ProjectID: f.ids.ProjectID, AgentID: f.agentID, ToolCallID: calls[0].ID, RuntimeLockID: work.RuntimeLock.ID, }, func(*executionstore.ToolCallReader) (executionstore.ToolCallCommand, error) { - return executionstore.CreateQuestionForToolCall(executionstore.CreateQuestionInteractionInput{Form: form}), nil + return executionstore.CreateQuestionForToolCall( + executionstore.CreateQuestionInteractionInput{Form: form}, + ), nil }) require.NoError(t, err) require.Equal(t, executionstore.ToolCallDispositionWaiting, question.Disposition) diff --git a/internal/e2e/service_e2e_deterministic_test.go b/internal/e2e/service_e2e_deterministic_test.go index ae36c1c32..42eb821d9 100644 --- a/internal/e2e/service_e2e_deterministic_test.go +++ b/internal/e2e/service_e2e_deterministic_test.go @@ -82,7 +82,9 @@ func TestServiceE2EDeterministicWorkerRunsModelTurn(t *testing.T) { err := env.db.QueryRow( ctx, `SELECT count(*) FROM agent_events event JOIN agents agent ON agent.id = event.agent_id JOIN content_blocks block ON block.agent_id = event.agent_id AND block.owner_model_output_id = event.model_output_id WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.event_kind = 'model_output' AND block.block_kind = 'text' AND block.text_content = $3`, - projectUUID, agentUUID, modelText, + projectUUID, + agentUUID, + modelText, ). Scan(&count) if err != nil { @@ -290,7 +292,9 @@ func TestServiceE2EDeterministicOpenRouterRunsChatCompletionsModelTurn(t *testin err := env.db.QueryRow( ctx, `SELECT count(*) FROM agent_events event JOIN agents agent ON agent.id = event.agent_id JOIN content_blocks block ON block.agent_id = event.agent_id AND block.owner_model_output_id = event.model_output_id WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.event_kind = 'model_output' AND block.block_kind = 'text' AND block.text_content = $3`, - projectUUID, agentUUID, modelText, + projectUUID, + agentUUID, + modelText, ). Scan(&count) if err != nil { @@ -380,12 +384,17 @@ func TestServiceE2EConfigChangeAffectsNextModelContext(t *testing.T) { defer openai.Close() env.startAPI(t, ctx) - project := env.bootstrapProjectViaAPIWithSource(t, ctx, "deterministic-config-change-runtime", strings.Join([]string{ - "instruction: Initial runtime instruction.", - "model:", - " provider_config: openai-prod", - " name: service-e2e-local", - }, "\n")+"\n") + project := env.bootstrapProjectViaAPIWithSource( + t, + ctx, + "deterministic-config-change-runtime", + strings.Join([]string{ + "instruction: Initial runtime instruction.", + "model:", + " provider_config: openai-prod", + " name: service-e2e-local", + }, "\n")+"\n", + ) agentID := project.createAgent(t, ctx) project.createInput(t, ctx, agentID, "first turn before config change") firstWorker := env.startWorker( @@ -401,13 +410,15 @@ func TestServiceE2EConfigChangeAffectsNextModelContext(t *testing.T) { err := env.db.QueryRow( ctx, `SELECT count(*) FROM agent_events event JOIN agents agent ON agent.id = event.agent_id JOIN content_blocks block ON block.agent_id = event.agent_id AND block.owner_model_output_id = event.model_output_id WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.event_kind = 'model_output' AND block.text_content = 'config change response 1'`, - projectUUID, agentUUID, + projectUUID, + agentUUID, ). Scan(&count) if err != nil { return false, err.Error() } - return count == 1, "first config-change runtime output not recorded; worker_logs=" + firstWorker.logExcerpt() + return count == 1, + "first config-change runtime output not recorded; worker_logs=" + firstWorker.logExcerpt() }) firstWorker.stop() @@ -432,13 +443,15 @@ func TestServiceE2EConfigChangeAffectsNextModelContext(t *testing.T) { err := env.db.QueryRow( ctx, `SELECT count(*) FROM agent_events event JOIN agents agent ON agent.id = event.agent_id JOIN content_blocks block ON block.agent_id = event.agent_id AND block.owner_model_output_id = event.model_output_id WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.event_kind = 'model_output' AND block.text_content = 'config change response 2'`, - projectUUID, agentUUID, + projectUUID, + agentUUID, ). Scan(&count) if err != nil { return false, err.Error() } - return count == 1, "second config-change runtime output not recorded; worker_logs=" + secondWorker.logExcerpt() + return count == 1, + "second config-change runtime output not recorded; worker_logs=" + secondWorker.logExcerpt() }) requestsMu.Lock() @@ -724,7 +737,9 @@ WHERE input.project_id = $1 err := env.db.QueryRow( ctx, `SELECT count(*) FROM agent_events event JOIN agents agent ON agent.id = event.agent_id JOIN content_blocks block ON block.agent_id = event.agent_id AND block.owner_model_output_id = event.model_output_id WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.event_kind = 'model_output' AND block.block_kind = 'text' AND block.text_content = $3`, - projectUUID, agentUUID, finalText, + projectUUID, + agentUUID, + finalText, ). Scan(&count) if err != nil { @@ -770,16 +785,12 @@ WHERE call.project_id = $1 var contextCount, distinctTurns int if err := env.db.QueryRow(ctx, ` SELECT count(*), - count(DISTINCT context_turn.turn_id), + count(DISTINCT context.turn_id), coalesce(string_agg( context.state || ':' || provider_config.api_format, ',' ORDER BY context.input_event_sequence, context.attempt_number, context.created_at, context.id ), '') FROM model_call_contexts context -JOIN model_call_context_turns context_turn - ON context_turn.project_id = context.project_id - AND context_turn.agent_id = context.agent_id - AND context_turn.model_call_context_id = context.id JOIN configured_model_revisions revision ON revision.org_id = context.org_id AND revision.id = context.configured_model_revision_id @@ -834,7 +845,8 @@ func TestServiceE2EDeterministicBacklogSteeringCancelAndQueuedContinuation(t *te if !strings.Contains(requestText, "steering priority message") { t.Errorf("first request did not admit steering input: %+v", body) } - if strings.Contains(requestText, "queued second message") || strings.Contains(requestText, "queued third message") || + if strings.Contains(requestText, "queued second message") || + strings.Contains(requestText, "queued third message") || strings.Contains(requestText, "queued canceled message") { t.Errorf("first request should not admit queued backlog while steering exists: %+v", body) } @@ -843,7 +855,8 @@ func TestServiceE2EDeterministicBacklogSteeringCancelAndQueuedContinuation(t *te if strings.Contains(requestText, "queued canceled message") { t.Errorf("second request included canceled backlog input: %+v", body) } - if !strings.Contains(requestText, "queued third message") || strings.Contains(requestText, "queued second message") { + if !strings.Contains(requestText, "queued third message") || + strings.Contains(requestText, "queued second message") { t.Errorf("second request should admit only the reordered front queued input: %+v", body) } writeOpenAIMessage(w, nil, "resp_queued_controls_first", "queued third turn complete") @@ -909,7 +922,8 @@ func TestServiceE2EDeterministicBacklogSteeringCancelAndQueuedContinuation(t *te http.StatusOK, ) backlogData := testutil.RequireType[[]any](t, backlog["data"]) - if len(backlogData) != 3 || testutil.RequireType[map[string]any](t, backlogData[0])["id"] != steeringInputID || + if len(backlogData) != 3 || + testutil.RequireType[map[string]any](t, backlogData[0])["id"] != steeringInputID || testutil.RequireType[map[string]any](t, backlogData[1])["id"] != thirdInputID || testutil.RequireType[map[string]any](t, backlogData[2])["id"] != secondInputID { t.Fatalf( @@ -1119,7 +1133,10 @@ func (e *serviceE2EEnvironment) bootstrapProjectViaAPIWithSourceAndModelOptions( ctx, http.MethodPost, projectPath+"/agent-profiles", - map[string]any{"name": "Deterministic Service E2E", "config": testutil.RequireType[string](t, config["id"])}, + map[string]any{ + "name": "Deterministic Service E2E", + "config": testutil.RequireType[string](t, config["id"]), + }, "idem-"+seed+"-agent-profile", adminToken, http.StatusCreated, @@ -1466,7 +1483,10 @@ func (p *deterministicProject) createInputWithDeliveryMode( p.adminToken, http.StatusCreated, ) - return testutil.RequireType[string](t, testutil.RequireType[map[string]any](t, created["agent_input"])["id"]) + return testutil.RequireType[string]( + t, + testutil.RequireType[map[string]any](t, created["agent_input"])["id"], + ) } func waitForServiceE2ECondition(t *testing.T, ctx context.Context, ready func() (bool, string)) { diff --git a/internal/e2e/service_e2e_durable_history_test.go b/internal/e2e/service_e2e_durable_history_test.go index 8a2a9304a..5aaad9e10 100644 --- a/internal/e2e/service_e2e_durable_history_test.go +++ b/internal/e2e/service_e2e_durable_history_test.go @@ -199,12 +199,8 @@ WHERE call.project_id = $1 t.Fatalf("query historical tool call: %v", err) } if err := env.db.QueryRow(ctx, ` -SELECT count(*), count(DISTINCT context_turn.turn_id) +SELECT count(*), count(DISTINCT context.turn_id) FROM model_call_contexts context -JOIN model_call_context_turns context_turn - ON context_turn.project_id = context.project_id - AND context_turn.agent_id = context.agent_id - AND context_turn.model_call_context_id = context.id WHERE context.project_id = $1 AND context.agent_id = $2 AND context.operation_kind = 'normal' @@ -395,12 +391,8 @@ WHERE context.project_id = $1 var contexts, distinctTurns, durableErrors int if err := env.db.QueryRow(ctx, ` -SELECT count(*), count(DISTINCT context_turn.turn_id) +SELECT count(*), count(DISTINCT context.turn_id) FROM model_call_contexts context -JOIN model_call_context_turns context_turn - ON context_turn.project_id = context.project_id - AND context_turn.agent_id = context.agent_id - AND context_turn.model_call_context_id = context.id WHERE context.project_id = $1 AND context.agent_id = $2 AND context.operation_kind = 'normal' @@ -592,12 +584,18 @@ func waitForServiceE2EAgentIdle( } if err := env.db.QueryRow( ctx, - `SELECT count(*) FROM agent_next_model_work($1,$2)`, + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND h.agent_id=$2 AND h.logical_ready_at IS NOT NULL`, projectUUID, agentUUID, ).Scan(&next); err != nil { return false, err.Error() } - return locks == 0 && wakeups == 0 && next == 0, fmt.Sprintf("locks=%d wakeups=%d next=%d", locks, wakeups, next) + return locks == 0 && wakeups == 0 && + next == 0, fmt.Sprintf( + "locks=%d wakeups=%d next=%d", + locks, + wakeups, + next, + ) }) } diff --git a/internal/harness/kernel/agent_executor.go b/internal/harness/kernel/agent_executor.go index 91276291e..cc4d1644c 100644 --- a/internal/harness/kernel/agent_executor.go +++ b/internal/harness/kernel/agent_executor.go @@ -36,18 +36,31 @@ type AgentExecutor struct { } func (e AgentExecutor) ExecuteModelWork(ctx context.Context, input ModelWorkExecution) error { + _, _, err := e.executeModelWork(ctx, input, nil) + return err +} + +func (e AgentExecutor) ExecuteModelWorkAndAdvance( + ctx context.Context, input ModelWorkExecution, advance ModelWorkAdvanceOptions, +) (executionstore.OwnedAgentWorkTransition, bool, error) { + return e.executeModelWork(ctx, input, &advance) +} + +func (e AgentExecutor) executeModelWork( + ctx context.Context, input ModelWorkExecution, advance *ModelWorkAdvanceOptions, +) (executionstore.OwnedAgentWorkTransition, bool, error) { if e.Store == nil { - return errors.New("kernel store is required") + return executionstore.OwnedAgentWorkTransition{}, false, errors.New("kernel store is required") } if err := validateModelWorkExecution(input); err != nil { - return err + return executionstore.OwnedAgentWorkTransition{}, false, err } if input.Now.IsZero() { input.Now = e.now() } builder := e.contextBuilder() if e.ModelResolver == nil { - return errors.New("kernel model resolver is required") + return executionstore.OwnedAgentWorkTransition{}, false, errors.New("kernel model resolver is required") } if input.Kind == executionstore.ModelWorkResume { @@ -58,13 +71,13 @@ func (e AgentExecutor) ExecuteModelWork(ctx context.Context, input ModelWorkExec input.ModelCallContextID, ) if loadErr != nil { - return loadErr + return executionstore.OwnedAgentWorkTransition{}, false, loadErr } if !found { - return errors.New("resume model context is missing") + return executionstore.OwnedAgentWorkTransition{}, false, errors.New("resume model context is missing") } if contextRow.OperationKind == executionstore.ModelCallOperationCompaction { - return e.resumeCompactionContext( + return executionstore.OwnedAgentWorkTransition{}, false, e.resumeCompactionContext( ctx, input, builder, @@ -74,13 +87,16 @@ func (e AgentExecutor) ExecuteModelWork(ctx context.Context, input ModelWorkExec } } - step, err := e.executeModelStep(ctx, input, builder, e.ModelResolver) + step, err := e.executeModelStepWithAdvance(ctx, input, builder, e.ModelResolver, advance) if err != nil { - return err + return executionstore.OwnedAgentWorkTransition{}, false, err + } + if step.Transition != nil { + return *step.Transition, true, nil } switch step.State { case modelStepWaiting, modelStepDone: - return nil + return executionstore.OwnedAgentWorkTransition{}, false, nil case modelStepToolUse: if _, err := e.recordToolCallSourceEvent( ctx, @@ -91,11 +107,11 @@ func (e AgentExecutor) ExecuteModelWork(ctx context.Context, input ModelWorkExec step.Bundle.ToolSpecs, step.StreamedToolCallIDs, ); err != nil { - return fmt.Errorf("record tool-call source event: %w", err) + return executionstore.OwnedAgentWorkTransition{}, false, fmt.Errorf("record tool-call source event: %w", err) } - return nil + return executionstore.OwnedAgentWorkTransition{}, false, nil default: - return fmt.Errorf("unsupported model step state %q", step.State) + return executionstore.OwnedAgentWorkTransition{}, false, fmt.Errorf("unsupported model step state %q", step.State) } } diff --git a/internal/harness/kernel/agent_executor_compaction_boundary_integration_test.go b/internal/harness/kernel/agent_executor_compaction_boundary_integration_test.go index 66a29d713..3fd8c0f42 100644 --- a/internal/harness/kernel/agent_executor_compaction_boundary_integration_test.go +++ b/internal/harness/kernel/agent_executor_compaction_boundary_integration_test.go @@ -70,12 +70,7 @@ func TestAgentExecutorRecoversInterruptedRetryCompaction(t *testing.T) { if err != nil { t.Fatalf("max event sequence before interrupted compaction: %v", err) } - snapshot, err := fixture.Store.Execution().CaptureAgentConfigForEventWatermark( - ctx, kernelTestProjectID, agentID, watermark, - ) - if err != nil { - t.Fatalf("capture config for interrupted compaction: %v", err) - } + recoveryModel := &sequenceKernelModel{ providerModelSlug: "kernel-test", capabilities: model.Capabilities{ @@ -84,25 +79,30 @@ func TestAgentExecutorRecoversInterruptedRetryCompaction(t *testing.T) { }, responses: []model.Response{ { - ID: "resp_recovered_compaction_summary", - Content: []model.ResponsePart{{Type: "text", Text: "Recovered summary for interrupted compaction."}}, + ID: "resp_recovered_compaction_summary", + Content: []model.ResponsePart{ + {Type: "text", Text: "Recovered summary for interrupted compaction."}, + }, StopReason: model.StopReasonEndTurn, }, { - ID: "resp_after_recovered_compaction", - Content: []model.ResponsePart{{Type: "text", Text: "continued after recovering interrupted compaction"}}, + ID: "resp_after_recovered_compaction", + Content: []model.ResponsePart{ + {Type: "text", Text: "continued after recovering interrupted compaction"}, + }, StopReason: model.StopReasonEndTurn, }, }, } - failedClaim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - RuntimeLockID: retryTurn.RuntimeLockID, - OpeningInputIDs: retryTurn.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: watermark, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: agentID, + RuntimeLockID: retryTurn.RuntimeLockID, + OpeningInputIDs: retryTurn.InputIDs, + }) + failedClaim := prepared1.Claim + if err != nil { t.Fatalf("claim context-window model call: %v", err) } @@ -195,7 +195,10 @@ func TestAgentExecutorRecoversInterruptedRetryCompaction(t *testing.T) { t.Fatalf("execute recovered interrupted compaction: %v", err) } if recoveryModel.respondedCount() != 1 { - t.Fatalf("recovery model prepared %d requests, want one compaction resend", recoveryModel.respondedCount()) + t.Fatalf( + "recovery model prepared %d requests, want one compaction resend", + recoveryModel.respondedCount(), + ) } parentContext, found, err := fixture.Store.Execution().GetModelCallContext( ctx, @@ -222,7 +225,10 @@ WHERE project_id = $1 AND input_event_sequence = $3 ORDER BY attempt_number DESC LIMIT 1 -`, kernelTestProjectID, agentID, failedClaim.Context.InputEventSequence).Scan(&recoveredCompactionContextID); err != nil { +`, + kernelTestProjectID, + agentID, + failedClaim.Context.InputEventSequence).Scan(&recoveredCompactionContextID); err != nil { t.Fatalf("load recovered compaction context id: %v", err) } reloadedCompaction, found, err := fixture.Store.Execution().GetModelCallContext( @@ -252,7 +258,8 @@ LIMIT 1 if err != nil { t.Fatalf("claim normal work after recovered compaction: %v", err) } - if !found || finalClaim.Kind != executionstore.AgentWorkModel || finalClaim.Model.ModelCallContextID != uuid.Nil { + if !found || finalClaim.Kind != executionstore.AgentWorkModel || + finalClaim.Model.ModelCallContextID != uuid.Nil { t.Fatalf("post-compaction work = %+v found=%v, want a fresh normal call", finalClaim, found) } finalTurn := modelWorkExecutionFromClaimForKernelTest(finalClaim, recoveryNow.Add(2*time.Second)) @@ -260,7 +267,10 @@ LIMIT 1 t.Fatalf("execute normal call after recovered compaction: %v", err) } if recoveryModel.respondedCount() != 2 { - t.Fatalf("recovery model prepared %d requests, want compaction and final call", recoveryModel.respondedCount()) + t.Fatalf( + "recovery model prepared %d requests, want compaction and final call", + recoveryModel.respondedCount(), + ) } var finalOutputs int if err := fixture.Pool.QueryRow( @@ -319,8 +329,10 @@ func TestAgentExecutorSteeringPreemptsCompactionCreatedFromInFlightOverflow(t *t StopReason: model.StopReasonContextWindow, }, { - ID: "resp_after_inflight_steering", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "used the steered request"}}, + ID: "resp_after_inflight_steering", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "used the steered request"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -331,10 +343,12 @@ func TestAgentExecutorSteeringPreemptsCompactionCreatedFromInFlightOverflow(t *t steeringInput, _, _, steeringErr = fixture.Store.Execution().CreateAgentContentInput( ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - Actor: kernelTestOmnaraActorParams(t, userID), - ContentBlocks: mustKernelJSON([]map[string]string{{"type": "text", "text": "use the newer request"}}), + ProjectID: kernelTestProjectID, + AgentID: agentID, + Actor: kernelTestOmnaraActorParams(t, userID), + ContentBlocks: mustKernelJSON( + []map[string]string{{"type": "text", "text": "use the newer request"}}, + ), DeliveryMode: executionstore.DeliveryModeSteering, IdempotencyKey: "steering-during-overflow", }, @@ -428,7 +442,10 @@ func TestAgentExecutorSteeringPreemptsCompactionCreatedFromInFlightOverflow(t *t } if modelClient.respondedCount() != 2 || !strings.Contains(string(modelClient.responded[1].ProviderRequest), "use the newer request") { - t.Fatalf("steered provider requests = %+v, want only the fresh request after overflow", modelClient.responded) + t.Fatalf( + "steered provider requests = %+v, want only the fresh request after overflow", + modelClient.responded, + ) } var oldState, freshState executionstore.ModelCallState var freshAttempt int @@ -497,7 +514,12 @@ func TestAgentExecutorConfigChangePreemptsSmallerCompactionCreatedFromTruncatedS nextConfig := fixture.kernelAgentConfigInput(t, ctx, "Kernel Test", "truncated-compaction-config-next") currentConfig := fixture.currentAgentConfig(t, ctx, agentID) currentRevisionID := currentRevisionIDForKernelConfig(t, ctx, fixture.Store, currentConfig) - nextRevisionID := currentRevisionIDForKernelConfiguredModelID(t, ctx, fixture.Store, nextConfig.ConfiguredModelID) + nextRevisionID := currentRevisionIDForKernelConfiguredModelID( + t, + ctx, + fixture.Store, + nextConfig.ConfiguredModelID, + ) var changeResult executionstore.ChangeAgentConfigResult var changeErr error @@ -518,8 +540,10 @@ func TestAgentExecutorConfigChangePreemptsSmallerCompactionCreatedFromTruncatedS StopReason: model.StopReasonMaxTokens, }, { - ID: "resp_smaller_compaction_should_not_send", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "stale compaction"}}, + ID: "resp_smaller_compaction_should_not_send", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "stale compaction"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -527,20 +551,23 @@ func TestAgentExecutorConfigChangePreemptsSmallerCompactionCreatedFromTruncatedS if response.ID != "resp_truncated_summary_before_config" { return } - changeResult, changeErr = fixture.Store.Execution().ChangeAgentConfig(ctx, executionstore.ChangeAgentConfigInput{ - CreateAgentConfigInput: nextConfig, - AgentID: agentID, - ActorType: identitystore.PrincipalTypeSystem, - Reason: "test_inflight_compaction_change", - IdempotencyKey: "config-during-truncated-compaction", - }) + changeResult, changeErr = fixture.Store.Execution(). + ChangeAgentConfig(ctx, executionstore.ChangeAgentConfigInput{ + CreateAgentConfigInput: nextConfig, + AgentID: agentID, + ActorType: identitystore.PrincipalTypeSystem, + Reason: "test_inflight_compaction_change", + IdempotencyKey: "config-during-truncated-compaction", + }) }, } newModel := &sequenceKernelModel{ providerModelSlug: "truncated-compaction-config-next", responses: []model.Response{{ - ID: "resp_after_truncated_compaction_config", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "used the changed config"}}, + ID: "resp_after_truncated_compaction_config", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "used the changed config"}, + }, StopReason: model.StopReasonEndTurn, }}, } @@ -584,7 +611,8 @@ func TestAgentExecutorConfigChangePreemptsSmallerCompactionCreatedFromTruncatedS oldModel.mu.Lock() remainingOldResponses := append([]model.Response(nil), oldModel.responses...) oldModel.mu.Unlock() - if len(remainingOldResponses) != 1 || remainingOldResponses[0].ID != "resp_smaller_compaction_should_not_send" { + if len(remainingOldResponses) != 1 || + remainingOldResponses[0].ID != "resp_smaller_compaction_should_not_send" { t.Fatalf("smaller stale compaction consumed a provider response: %+v", remainingOldResponses) } if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( @@ -715,13 +743,17 @@ func TestAgentExecutorSteeringStartsFreshFrontierDuringCompactionRetry(t *testin }, responses: []model.Response{ { - ID: "resp_context_window_before_new_input", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "context window exceeded"}}, + ID: "resp_context_window_before_new_input", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "context window exceeded"}, + }, StopReason: model.StopReasonContextWindow, }, { - ID: "resp_after_compaction_retry_supersession", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "handled the newer input"}}, + ID: "resp_after_compaction_retry_supersession", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "handled the newer input"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -745,7 +777,10 @@ func TestAgentExecutorSteeringStartsFreshFrontierDuringCompactionRetry(t *testin t.Fatalf("execute overflow and retrying compaction: %v", err) } if retryModel.respondedCount() != 2 { - t.Fatalf("prepared requests = %d, want normal overflow and compaction retry", retryModel.respondedCount()) + t.Fatalf( + "prepared requests = %d, want normal overflow and compaction retry", + retryModel.respondedCount(), + ) } var blockedContextID, compactionContextID uuid.UUID var compactionRetryAt time.Time @@ -794,8 +829,13 @@ func TestAgentExecutorSteeringStartsFreshFrontierDuringCompactionRetry(t *testin ProjectID: kernelTestProjectID, AgentID: agentID, InputEventSequence: retryingCompactionContext.InputEventSequence, - EventSequenceStart: compactionSourceStartForKernelTest(t, ctx, fixture.Store, retryingCompactionContext), - EventSequenceEnd: *retryingCompactionContext.SourceEventSequenceEnd, + EventSequenceStart: compactionSourceStartForKernelTest( + t, + ctx, + fixture.Store, + retryingCompactionContext, + ), + EventSequenceEnd: *retryingCompactionContext.SourceEventSequenceEnd, }, TurnID: overflowTurn.TurnID, OpeningInputIDs: overflowTurn.InputIDs, diff --git a/internal/harness/kernel/agent_executor_compaction_progressive_integration_test.go b/internal/harness/kernel/agent_executor_compaction_progressive_integration_test.go index f115d3eb6..ec34c06c3 100644 --- a/internal/harness/kernel/agent_executor_compaction_progressive_integration_test.go +++ b/internal/harness/kernel/agent_executor_compaction_progressive_integration_test.go @@ -124,8 +124,10 @@ func TestAgentExecutorProgressiveCompactionCompletesWithoutReexpandingSource(t * StopReason: model.StopReasonEndTurn, }, { - ID: "resp_after_progressive_compaction", - Content: []model.ResponsePart{{Type: "text", Text: "continued after progressive compaction"}}, + ID: "resp_after_progressive_compaction", + Content: []model.ResponsePart{ + {Type: "text", Text: "continued after progressive compaction"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -256,7 +258,13 @@ func TestAgentExecutorProgressiveCompactionCompletesWithoutReexpandingSource(t * test.frontier, ) if err != nil || depth != test.wantDepth { - t.Fatalf("%s checkpoint lineage depth = %d, want %d (err=%v)", test.name, depth, test.wantDepth, err) + t.Fatalf( + "%s checkpoint lineage depth = %d, want %d (err=%v)", + test.name, + depth, + test.wantDepth, + err, + ) } } semanticFrontier, err := fixture.Store.Execution().MaxEventSequence(ctx, kernelTestProjectID, agentID) @@ -394,7 +402,11 @@ func TestProgressiveCompactionExhaustionPublishesOneParentError(t *testing.T) { t.Fatalf("load progressive exhaustion watermark: %v", err) } if watermark != turn.OpeningEventSequence || watermark != 10 { - t.Fatalf("progressive exhaustion watermark/opening = %d/%d, want 10", watermark, turn.OpeningEventSequence) + t.Fatalf( + "progressive exhaustion watermark/opening = %d/%d, want 10", + watermark, + turn.OpeningEventSequence, + ) } compactionModel := &sequenceKernelModel{ @@ -422,7 +434,11 @@ func TestProgressiveCompactionExhaustionPublishesOneParentError(t *testing.T) { Store: compaction.NewStore(fixture.Store.Execution()), Resolver: resolver, ContextBuilder: modelcontext.Builder{ - Store: modelcontext.NewStore(fixture.Store.Execution(), fixture.Store.Artifacts(), fixture.Store.Integrations()), + Store: modelcontext.NewStore( + fixture.Store.Execution(), + fixture.Store.Artifacts(), + fixture.Store.Integrations(), + ), }, Now: func() time.Time { return runNow }, } @@ -430,23 +446,16 @@ func TestProgressiveCompactionExhaustionPublishesOneParentError(t *testing.T) { frontier, sourceEventSequenceEnd int64, ) (executionstore.ModelCallClaim, executionstore.ModelCallClaim) { t.Helper() - snapshot, err := fixture.Store.Execution().CaptureAgentConfigForEventWatermark( - ctx, - kernelTestProjectID, - agentID, - frontier, - ) - if err != nil { - t.Fatalf("capture progressive exhaustion parent snapshot: %v", err) - } - parent, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - RuntimeLockID: turn.RuntimeLockID, - OpeningInputIDs: turn.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: frontier, - }) + + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: agentID, + RuntimeLockID: turn.RuntimeLockID, + OpeningInputIDs: turn.InputIDs, + }) + parent := prepared1.Claim + if err != nil { t.Fatalf("claim progressive exhaustion parent: %v", err) } diff --git a/internal/harness/kernel/agent_executor_compaction_provider_integration_test.go b/internal/harness/kernel/agent_executor_compaction_provider_integration_test.go index 9ccbb0203..e5716ffb8 100644 --- a/internal/harness/kernel/agent_executor_compaction_provider_integration_test.go +++ b/internal/harness/kernel/agent_executor_compaction_provider_integration_test.go @@ -32,13 +32,17 @@ func TestAgentExecutorCarriesReplayRejectionThroughCompaction(t *testing.T) { providerModelSlug: "replay-compaction-model", responses: []model.Response{ { - ID: "resp-replay-compaction-padding", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "disposable old history accepted"}}, + ID: "resp-replay-compaction-padding", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "disposable old history accepted"}, + }, StopReason: model.StopReasonEndTurn, }, { - ID: "resp-replay-compaction-seed", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "recent history that must remain exact"}}, + ID: "resp-replay-compaction-seed", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "recent history that must remain exact"}, + }, StopReason: model.StopReasonEndTurn, ProviderReplay: json.RawMessage( `[{"type":"reasoning","id":"rs_replay_compaction","encrypted_content":"large opaque replay"},` + @@ -125,14 +129,19 @@ func TestAgentExecutorCarriesReplayRejectionThroughCompaction(t *testing.T) { }}, responses: []model.Response{ { - ID: "resp-replay-compaction-overflow", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "context window exceeded"}}, + ID: "resp-replay-compaction-overflow", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "context window exceeded"}, + }, StopReason: model.StopReasonContextWindow, }, { ID: "resp-replay-compaction-summary", Content: []model.ResponsePart{ - {Type: model.ResponsePartTypeText, Text: "The earlier replay-backed history was preserved."}, + { + Type: model.ResponsePartTypeText, + Text: "The earlier replay-backed history was preserved.", + }, }, StopReason: model.StopReasonEndTurn, }, @@ -149,8 +158,10 @@ func TestAgentExecutorCarriesReplayRejectionThroughCompaction(t *testing.T) { ), }, { - ID: "resp-replay-compaction-later", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "new replay remained eligible"}}, + ID: "resp-replay-compaction-later", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "new replay remained eligible"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -256,7 +267,10 @@ func TestAgentExecutorCarriesReplayRejectionThroughCompaction(t *testing.T) { } } if checkpointProjections < 3 { - t.Fatalf("checkpoint projections = %d, want planner, validator, and continuation", checkpointProjections) + t.Fatalf( + "checkpoint projections = %d, want planner, validator, and continuation", + checkpointProjections, + ) } var checkpoints int @@ -352,8 +366,10 @@ func TestAgentExecutorCompactsAndRetriesAfterProviderContextWindow(t *testing.T) Usage: model.Usage{InputTokens: 128000}, }, { - ID: "resp_summary", - Content: []model.ResponsePart{{Type: "text", Text: "The earlier turn established seed history."}}, + ID: "resp_summary", + Content: []model.ResponsePart{ + {Type: "text", Text: "The earlier turn established seed history."}, + }, StopReason: model.StopReasonEndTurn, Usage: model.Usage{InputTokens: 20, OutputTokens: 8}, }, @@ -392,7 +408,9 @@ func TestAgentExecutorCompactsAndRetriesAfterProviderContextWindow(t *testing.T) ) } if len(retryModel.responded[0].ToolSpecs) == 0 { - t.Fatal("normal request omitted the configured tool needed to make the compaction assertion meaningful") + t.Fatal( + "normal request omitted the configured tool needed to make the compaction assertion meaningful", + ) } if len(retryModel.responded[1].ToolSpecs) != 0 { t.Fatalf("compaction request exposed tools: %+v", retryModel.responded[1].ToolSpecs) @@ -441,8 +459,13 @@ LIMIT 1 ProjectID: kernelTestProjectID, AgentID: agentID, InputEventSequence: completedCompactionContext.InputEventSequence, - EventSequenceStart: compactionSourceStartForKernelTest(t, ctx, fixture.Store, completedCompactionContext), - EventSequenceEnd: *completedCompactionContext.SourceEventSequenceEnd, + EventSequenceStart: compactionSourceStartForKernelTest( + t, + ctx, + fixture.Store, + completedCompactionContext, + ), + EventSequenceEnd: *completedCompactionContext.SourceEventSequenceEnd, }, TurnID: retryTurn.TurnID, OpeningInputIDs: retryTurn.InputIDs, @@ -457,7 +480,10 @@ LIMIT 1 t.Fatalf("replayed compaction = %+v, want adopted checkpoint", replayedCompaction) } if retryModel.respondedCount() != 2 { - t.Fatalf("completed compaction replay made a provider request; prepared=%d", retryModel.respondedCount()) + t.Fatalf( + "completed compaction replay made a provider request; prepared=%d", + retryModel.respondedCount(), + ) } finalTurn := continueTurnOnNewLeaseForKernelTest( t, @@ -521,7 +547,10 @@ WHERE agent.project_id = $1 AND checkpoint.agent_id = $2 AND mcc.operation_kind = 'normal' AND mcc.state = 'succeeded' AND mcc.input_event_sequence >= checkpoint_event.sequence - AND context_turn.turn_id = $3`, kernelTestProjectID, agentID, retryTurn.TurnID).Scan(&retryContexts); err != nil { + AND context_turn.turn_id = $3`, + kernelTestProjectID, + agentID, + retryTurn.TurnID).Scan(&retryContexts); err != nil { t.Fatalf("count retry context checkpoint refs: %v", err) } if retryContexts != 1 { @@ -549,8 +578,10 @@ func TestAgentExecutorReplaysOverflowWhenPlanningIsInterruptedBeforeHandoff(t *t seedModel := &sequenceKernelModel{ providerModelSlug: "kernel-test", responses: []model.Response{{ - ID: "resp_interrupted_compaction_seed", - Content: []model.ResponsePart{{Type: "text", Text: "closed history before interrupted compaction"}}, + ID: "resp_interrupted_compaction_seed", + Content: []model.ResponsePart{ + {Type: "text", Text: "closed history before interrupted compaction"}, + }, StopReason: model.StopReasonEndTurn, }}, } @@ -589,8 +620,10 @@ func TestAgentExecutorReplaysOverflowWhenPlanningIsInterruptedBeforeHandoff(t *t }, responses: []model.Response{ { - ID: "resp_overflow_before_interrupted_handoff", - Content: []model.ResponsePart{{Type: "text", Text: "context window exceeded before handoff"}}, + ID: "resp_overflow_before_interrupted_handoff", + Content: []model.ResponsePart{ + {Type: "text", Text: "context window exceeded before handoff"}, + }, StopReason: model.StopReasonContextWindow, Usage: model.Usage{InputTokens: 128000}, }, @@ -608,8 +641,10 @@ func TestAgentExecutorReplaysOverflowWhenPlanningIsInterruptedBeforeHandoff(t *t StopReason: model.StopReasonEndTurn, }, { - ID: "resp_final_after_ambiguous_replay", - Content: []model.ResponsePart{{Type: "text", Text: "continued after interrupted compaction recovery"}}, + ID: "resp_final_after_ambiguous_replay", + Content: []model.ResponsePart{ + {Type: "text", Text: "continued after interrupted compaction recovery"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -723,7 +758,10 @@ WHERE agent.project_id = $1 AND checkpoint.agent_id = $2 t.Fatalf("execute ambiguous overflow replay and compaction: %v", err) } if modelClient.respondedCount() != 3 { - t.Fatalf("requests through replay compaction = %d, want overflow/overflow/summary", modelClient.respondedCount()) + t.Fatalf( + "requests through replay compaction = %d, want overflow/overflow/summary", + modelClient.respondedCount(), + ) } finalNow := claimAt.Add(time.Second) @@ -733,7 +771,10 @@ WHERE agent.project_id = $1 AND checkpoint.agent_id = $2 t.Fatalf("execute continuation after interrupted compaction recovery: %v", err) } if modelClient.respondedCount() != 4 { - t.Fatalf("total requests after interrupted compaction recovery = %d, want four", modelClient.respondedCount()) + t.Fatalf( + "total requests after interrupted compaction recovery = %d, want four", + modelClient.respondedCount(), + ) } if summaryRequest := string(modelClient.responded[2].ProviderRequest); !strings.Contains( summaryRequest, @@ -906,7 +947,10 @@ func TestAgentExecutorStopsManagedCompactionRetryAfterAdmissionCloses(t *testing t.Fatalf("execute retryable managed compaction: %v", err) } if journey.model.respondedCount() != 2 { - t.Fatalf("provider requests through retryable compaction = %d, want 2", journey.model.respondedCount()) + t.Fatalf( + "provider requests through retryable compaction = %d, want 2", + journey.model.respondedCount(), + ) } journey.fixture.setManagedWorkAdmission(t, ctx, false) @@ -922,7 +966,10 @@ func TestAgentExecutorStopsManagedCompactionRetryAfterAdmissionCloses(t *testing t.Fatalf("deny managed compaction retry: %v", err) } if journey.model.respondedCount() != 2 { - t.Fatalf("provider requests after denied compaction retry = %d, want 2", journey.model.respondedCount()) + t.Fatalf( + "provider requests after denied compaction retry = %d, want 2", + journey.model.respondedCount(), + ) } assertDurableModelErrorForKernelTest( t, @@ -959,7 +1006,7 @@ WHERE project_id = $1 func TestManagedCompactionSourceReplacementStopsAfterAdmissionCloses(t *testing.T) { ctx := context.Background() journey := newManagedCompactionAdmissionJourney(t, ctx, "") - frontier, err := journey.fixture.Store.Execution().MaxEventSequence( + _, err := journey.fixture.Store.Execution().MaxEventSequence( ctx, kernelTestProjectID, journey.agentID, @@ -967,26 +1014,18 @@ func TestManagedCompactionSourceReplacementStopsAfterAdmissionCloses(t *testing. if err != nil { t.Fatalf("load managed replacement frontier: %v", err) } - snapshot, err := journey.fixture.Store.Execution().CaptureAgentConfigForEventWatermark( - ctx, - kernelTestProjectID, - journey.agentID, - frontier, - ) - if err != nil { - t.Fatalf("capture managed replacement config: %v", err) - } - parent, err := journey.fixture.Store.Execution().ClaimNormalModelCall( + + prepared1, err := journey.fixture.Store.Execution().PrepareNormalModelCall( ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: journey.agentID, - RuntimeLockID: journey.turn.RuntimeLockID, - OpeningInputIDs: journey.turn.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: frontier, + executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: journey.agentID, + RuntimeLockID: journey.turn.RuntimeLockID, + OpeningInputIDs: journey.turn.InputIDs, }, ) + parent := prepared1.Claim + if err != nil { t.Fatalf("claim managed replacement parent: %v", err) } @@ -1154,8 +1193,10 @@ func newManagedCompactionAdmissionJourney( seedModel := &sequenceKernelModel{ providerModelSlug: "managed-compaction-model", responses: []model.Response{{ - ID: "resp_managed_seed", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "managed seed accepted"}}, + ID: "resp_managed_seed", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "managed seed accepted"}, + }, StopReason: model.StopReasonEndTurn, }}, } @@ -1183,13 +1224,17 @@ func newManagedCompactionAdmissionJourney( }, responses: []model.Response{ { - ID: "resp_managed_context_window", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "context window exceeded"}}, + ID: "resp_managed_context_window", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "context window exceeded"}, + }, StopReason: model.StopReasonContextWindow, }, { - ID: "resp_managed_summary", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "Managed history summary."}}, + ID: "resp_managed_summary", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "Managed history summary."}, + }, StopReason: model.StopReasonEndTurn, }, { @@ -1277,24 +1322,17 @@ func TestCompactionExhaustsMalformedResponsesWithoutPersistingUnsafeEvidence(t * if err != nil { t.Fatalf("load malformed compaction frontier: %v", err) } - snapshot, err := fixture.Store.Execution().CaptureAgentConfigForEventWatermark( - ctx, - kernelTestProjectID, - agentID, - frontier, - ) - if err != nil { - t.Fatalf("capture malformed compaction config: %v", err) - } + runNow := fixture.Now.Add(4 * time.Second) - parent, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - RuntimeLockID: turn.RuntimeLockID, - OpeningInputIDs: turn.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: frontier, - }) + prepared2, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: agentID, + RuntimeLockID: turn.RuntimeLockID, + OpeningInputIDs: turn.InputIDs, + }) + parent := prepared2.Claim + if err != nil { t.Fatalf("claim malformed compaction parent: %v", err) } @@ -1331,8 +1369,10 @@ func TestCompactionExhaustsMalformedResponsesWithoutPersistingUnsafeEvidence(t * malformedResponses[index] = model.Response{ ID: "resp\x00malformed_compaction_" + strconv.Itoa(index+1), ServedProviderModelSlug: "served\x00malformed_compaction", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "complete summary"}}, - StopReason: model.StopReasonEndTurn, + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "complete summary"}, + }, + StopReason: model.StopReasonEndTurn, } } compactionModel := &sequenceKernelModel{ @@ -1397,22 +1437,40 @@ func TestCompactionExhaustsMalformedResponsesWithoutPersistingUnsafeEvidence(t * } if result.State != wantState || contextRecord.AttemptNumber != attemptNumber || contextRecord.State != executionstore.ModelCallContextFailed || - contextRecord.RecoveryKind != wantRecoveryKind || contextRecord.ErrorCode != "malformed_success_response" || + contextRecord.RecoveryKind != wantRecoveryKind || + contextRecord.ErrorCode != "malformed_success_response" || contextRecord.ProviderResponseID != "" || !kernelModelCallOutcomeAmbiguous(t, contextRecord.ErrorDetails) { - t.Fatalf("malformed compaction attempt %d result=%+v context=%+v", attemptNumber, result, contextRecord) + t.Fatalf( + "malformed compaction attempt %d result=%+v context=%+v", + attemptNumber, + result, + contextRecord, + ) } if attemptNumber < maxAttempts { - if result.RetryAt == nil || contextRecord.RetryAt == nil || !result.RetryAt.Equal(*contextRecord.RetryAt) { - t.Fatalf("malformed compaction retry %d result=%+v context=%+v", attemptNumber, result, contextRecord) + if result.RetryAt == nil || contextRecord.RetryAt == nil || + !result.RetryAt.Equal(*contextRecord.RetryAt) { + t.Fatalf( + "malformed compaction retry %d result=%+v context=%+v", + attemptNumber, + result, + contextRecord, + ) } currentNow = *result.RetryAt } else if result.RetryAt != nil || contextRecord.RetryAt != nil { - t.Fatalf("terminal malformed compaction retained retry time: result=%+v context=%+v", result, contextRecord) + t.Fatalf("terminal malformed compaction retained retry time: result=%+v context=%+v", + result, + contextRecord) } } if compactionModel.respondedCount() != maxAttempts { - t.Fatalf("malformed compaction provider calls = %d, want %d", compactionModel.respondedCount(), maxAttempts) + t.Fatalf( + "malformed compaction provider calls = %d, want %d", + compactionModel.respondedCount(), + maxAttempts, + ) } var checkpoints int if err := fixture.Pool.QueryRow(ctx, ` @@ -1465,8 +1523,8 @@ GROUP BY context.state, output.stop_reason`, if err := fixture.Pool.QueryRow(ctx, ` SELECT (SELECT count(*)::integer FROM agent_wakeups wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.project_id = $1 AND wake.agent_id = $2), - (SELECT count(*)::integer FROM agent_continuable_model_contexts($1, $2) - WHERE model_call_context_id = $3 AND NOT has_later_semantic_event)`, + (SELECT count(*)::integer FROM agent_execution_state h + WHERE h.agent_id=$2 AND h.normal_context_id=$3 AND h.turn_continuable)`, kernelTestProjectID, agentID, result.ModelCallContextID).Scan(&wakeups, &resumable); err != nil { t.Fatalf("load terminal compaction continuation state: %v", err) } @@ -1478,8 +1536,16 @@ SELECT (SELECT count(*)::integer FROM agent_wakeups wake JOIN agents agent ON ag wakeups != 0 || resumable != 0 { t.Fatalf( "terminal malformed compaction state context=%s producer=%s contexts=%d output=%s/%d block=%s/%q/%d wakeups=%d resumable=%d", - compactionState, producingState, contexts, outputStopReason, outputCount, - blockKind, blockText, blockCount, wakeups, resumable, + compactionState, + producingState, + contexts, + outputStopReason, + outputCount, + blockKind, + blockText, + blockCount, + wakeups, + resumable, ) } } @@ -1614,7 +1680,10 @@ func TestAgentExecutorRecordsErrorWhenModelGrantDisappearsBeforeCompaction(t *te remainingResponses := append([]model.Response(nil), retryModel.responses...) retryModel.mu.Unlock() if len(remainingResponses) != 1 || remainingResponses[0].ID != "resp_compaction_should_not_send" { - t.Fatalf("compaction provider response was consumed after grant revoke; remaining=%+v", remainingResponses) + t.Fatalf( + "compaction provider response was consumed after grant revoke; remaining=%+v", + remainingResponses, + ) } var errorOutputs int if err := fixture.Pool.QueryRow(ctx, ` @@ -1696,7 +1765,8 @@ func TestAgentExecutorRecordsErrorWhenModelGrantDisappearsBeforeCompaction(t *te ); err != nil { t.Fatalf("load terminal compaction lineage: %v", err) } - if failedNormalContexts != 1 || failedCompactionContexts != 1 || normalContexts != 1 || compactionContexts != 1 { + if failedNormalContexts != 1 || failedCompactionContexts != 1 || normalContexts != 1 || + compactionContexts != 1 { t.Fatalf( "terminal compaction lineage failed normal/compaction=%d/%d context rows=%d/%d, want 1/1 and 1/1", failedNormalContexts, diff --git a/internal/harness/kernel/agent_executor_context.go b/internal/harness/kernel/agent_executor_context.go index f96d4aea9..39b9ad912 100644 --- a/internal/harness/kernel/agent_executor_context.go +++ b/internal/harness/kernel/agent_executor_context.go @@ -3,6 +3,7 @@ package kernel import ( "context" "errors" + "time" "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/agentconfig" @@ -34,6 +35,7 @@ const ( ) type modelStep struct { + Transition *executionstore.OwnedAgentWorkTransition State modelStepState Context executionstore.ModelCallContextRecord Bundle modelcontext.Bundle @@ -43,13 +45,12 @@ type modelStep struct { StreamedToolCallIDs map[string]uuid.UUID } -func (e AgentExecutor) executeModelStep( - ctx context.Context, - input ModelWorkExecution, - builder modelcontext.Builder, - resolver model.Resolver, +func (e AgentExecutor) executeModelStepWithAdvance( + ctx context.Context, input ModelWorkExecution, builder modelcontext.Builder, resolver model.Resolver, + advance *ModelWorkAdvanceOptions, ) (modelStep, error) { var snapshot executionstore.AgentConfigSnapshotRecord + var contextData *executionstore.ModelContextData var claim executionstore.ModelCallClaim var err error if input.Kind == executionstore.ModelWorkResume { @@ -80,23 +81,28 @@ func (e AgentExecutor) executeModelStep( } } } else { - snapshot, err = e.Store.Execution().CaptureAgentConfigForModelContext(ctx, input.ProjectID, input.AgentID) - if err == nil { - claim, err = e.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ + var prepared executionstore.PreparedNormalModelCall + if input.Prepared != nil { + prepared = *input.Prepared + } else { + prepared, err = e.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ ProjectID: input.ProjectID, AgentID: input.AgentID, RuntimeLockID: input.RuntimeLockID, OpeningInputIDs: input.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: snapshot.InputEventSequence, SourceModelCallContextID: input.SourceModelCallContextID, SourceModelOutputID: input.SourceModelOutputID, }) } + claim, snapshot, contextData = prepared.Claim, prepared.Snapshot, prepared.ContextData } if err != nil { return modelStep{}, err } + if input.Prepared != nil && (claim.Context.ProjectID != input.ProjectID || claim.Context.AgentID != input.AgentID || + claim.Context.RuntimeLockID != input.RuntimeLockID || claim.Context.AgentConfigID != snapshot.AgentConfig.ID) { + return modelStep{}, storeerr.ErrStateTransitionConflict + } if !claim.Claimed { if claim.Created && claim.Context.ErrorCode == storeerr.ManagedWorkAdmissionDeniedCode { @@ -185,6 +191,7 @@ func (e AgentExecutor) executeModelStep( OpeningInputIDs: input.InputIDs, Now: input.Now, AgentConfigSnapshot: &snapshot, + ContextData: contextData, MediaProjector: model.MediaProjectorForClient(client), }) if errors.Is(err, modelcontext.ErrOpeningMediaBudgetExceeded) { @@ -220,14 +227,21 @@ func (e AgentExecutor) executeModelStep( return e.recordNormalFailure(ctx, input, claim, resolved, cause, false, model.Response{}) } - policy, err := modelretry.RequestPolicyForModelCall( - ctx, - e.Store.Execution(), - input.ProjectID, - input.AgentID, - claim.Context.ID, - model.RequestPolicyFromCapabilities(capabilities), - ) + policy := model.RequestPolicyFromCapabilities(capabilities) + if contextData != nil { + policy.ProviderReplayCutoffEventSequence = max(policy.ProviderReplayCutoffEventSequence, + contextData.ProviderReplayCutoffEventSequence) + } else { + policy, err = modelretry.RequestPolicyForModelCall( + ctx, + e.Store.Execution(), + input.ProjectID, + input.AgentID, + claim.Context.ID, + model.RequestPolicyFromCapabilities(capabilities), + ) + } + if err != nil { return e.recordNormalPreSendFailure( ctx, input, claim, resolved, err, @@ -332,6 +346,26 @@ func (e AgentExecutor) executeModelStep( Resolved: resolved, StreamedToolCallIDs: streamSink.ToolCallIDs(), } + if advance != nil && time.Now().Before(advance.Deadline) && + invalidModelResponse(modelErrorSourceForClient(client), step.Envelope.Normalized.StopReason, + model.ToolCallsFromEnvelope(step.Envelope)) == nil { + transition, err := e.Store.Execution().AcceptModelOutputAndAdvance( + ctx, executionstore.AcceptModelOutputAndAdvanceInput{ + Output: executionstore.RecordModelOutputAndCompleteContextInput{ + ProjectID: input.ProjectID, AgentID: input.AgentID, RuntimeLockID: input.RuntimeLockID, + ModelCallContextID: claim.Context.ID, ProviderRequestID: step.Response.ProviderRequestID, + ProviderResponse: step.Envelope, + }, + ToolCallBindings: toolCallBindings(step.Envelope, step.Bundle.ToolSpecs, step.StreamedToolCallIDs), + AllowModelWork: advance.AllowModelWork, PrepareModel: true, + }, + ) + if err != nil { + return modelStep{}, err + } + step.State, step.Transition = modelStepDone, &transition + return step, nil + } return e.finishModelResponse(ctx, input, step) } diff --git a/internal/harness/kernel/agent_executor_model_context.go b/internal/harness/kernel/agent_executor_model_context.go index 88232e4b7..5f6d34ec7 100644 --- a/internal/harness/kernel/agent_executor_model_context.go +++ b/internal/harness/kernel/agent_executor_model_context.go @@ -76,6 +76,15 @@ func (e AgentExecutor) modelContextToolRuntime( if !found { return nil, errors.New("model context agent config is missing") } + return e.toolRuntimeFromConfig(ctx, projectID, agentID, config, now) +} + +func (e AgentExecutor) toolRuntimeFromConfig( + ctx context.Context, + projectID, agentID uuid.UUID, + config executionstore.AgentConfigRecord, + now time.Time, +) ([]modelcontext.ToolSpec, error) { contract, err := agentconfig.RuntimeContractFromCompiled( config.CompiledDefinition, config.EffectiveDefinitionHash, diff --git a/internal/harness/kernel/agent_executor_model_context_integration_test.go b/internal/harness/kernel/agent_executor_model_context_integration_test.go index 875712bb0..a91056333 100644 --- a/internal/harness/kernel/agent_executor_model_context_integration_test.go +++ b/internal/harness/kernel/agent_executor_model_context_integration_test.go @@ -78,15 +78,15 @@ model: if err != nil { t.Fatalf("capture agent config: %v", err) } - modelClaim, err := fixture.Store.Execution(). - ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: turn.RuntimeLockID, - OpeningInputIDs: turn.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: turn.OpeningEventSequence, + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: launch.Agent.ID, + RuntimeLockID: turn.RuntimeLockID, + OpeningInputIDs: turn.InputIDs, }) + modelClaim := prepared1.Claim + if err != nil { t.Fatalf("claim model call context: %v", err) } @@ -625,7 +625,8 @@ func (f kernelFixture) admitSlackContentInputTurn( require.NoError(t, err) require.Equal(t, executionstore.InteractionSelection{AutoSelect: true}, selection, "queued Slack content must not select a destination before turn admission") - claim, found, err := f.Store.Execution().ClaimNextAgentWork(ctx, kernelTestClaimInput(now.Add(2*time.Millisecond))) + claim, found, err := f.Store.Execution(). + ClaimNextAgentWork(ctx, kernelTestClaimInput(now.Add(2*time.Millisecond))) require.NoError(t, err) require.True(t, found) require.Equal(t, executionstore.AgentWorkModel, claim.Kind) diff --git a/internal/harness/kernel/agent_executor_model_request_integration_test.go b/internal/harness/kernel/agent_executor_model_request_integration_test.go index 8567e1e9a..6ef73584b 100644 --- a/internal/harness/kernel/agent_executor_model_request_integration_test.go +++ b/internal/harness/kernel/agent_executor_model_request_integration_test.go @@ -169,7 +169,9 @@ WHERE project_id = $1 AND agent_id = $2 AND input_event_sequence = $3 t.Fatalf("count denied managed model contexts: %v", err) } if contexts != 1 || len(resolver.selections) != 0 { - t.Fatalf("denied managed model contexts/resolver calls = %d/%d, want 1/0", contexts, len(resolver.selections)) + t.Fatalf("denied managed model contexts/resolver calls = %d/%d, want 1/0", + contexts, + len(resolver.selections)) } if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( ctx, @@ -550,7 +552,8 @@ model: t.Fatalf("resolver selections = %d, want 1", len(resolver.selections)) } selection := resolver.selections[0] - if selection.Overrides.DefaultMaxOutputTokens == nil || *selection.Overrides.DefaultMaxOutputTokens != 1234 || + if selection.Overrides.DefaultMaxOutputTokens == nil || + *selection.Overrides.DefaultMaxOutputTokens != 1234 || selection.Overrides.CacheRetention != string(model.CacheRetentionShort) { t.Fatalf("selection overrides = %+v", selection.Overrides) } @@ -681,7 +684,9 @@ func TestAgentExecutorCarriesDurableProviderReplayIntoNextTurn(t *testing.T) { } bundle, err := (modelcontext.Builder{ - Store: modelcontext.NewStore(fixture.Store.Execution(), fixture.Store.Artifacts(), fixture.Store.Integrations()), + Store: modelcontext.NewStore(fixture.Store.Execution(), + fixture.Store.Artifacts(), + fixture.Store.Integrations()), }).Build( ctx, modelcontext.BuildInput{ @@ -1095,6 +1100,7 @@ FROM model_call_contexts context WHERE context.project_id = $1 AND context.agent_id = $2 AND context.input_event_sequence = $3 + AND context.operation_kind='normal' `, kernelTestProjectID, agentID, turn.OpeningEventSequence).Scan( &state, &recoveryKind, @@ -1256,6 +1262,7 @@ FROM model_call_contexts context WHERE context.project_id = $1 AND context.agent_id = $2 AND context.input_event_sequence = $3 + AND context.operation_kind = 'normal' `, kernelTestProjectID, agentID, secondTurn.OpeningEventSequence).Scan( &state, &recoveryKind, @@ -1263,7 +1270,8 @@ WHERE context.project_id = $1 ); err != nil { t.Fatalf("load compactable serialized overflow attempt: %v", err) } - if state != executionstore.ModelCallContextFailed || recoveryKind != executionstore.ModelCallRecoveryCompact || + if state != executionstore.ModelCallContextFailed || + recoveryKind != executionstore.ModelCallRecoveryCompact || errorCode != "configured_input_budget_exceeded" { t.Fatalf( "compactable serialized overflow attempt = %q/%q/%q", diff --git a/internal/harness/kernel/agent_executor_output_truncation_integration_test.go b/internal/harness/kernel/agent_executor_output_truncation_integration_test.go index 676096ac7..7db152cae 100644 --- a/internal/harness/kernel/agent_executor_output_truncation_integration_test.go +++ b/internal/harness/kernel/agent_executor_output_truncation_integration_test.go @@ -35,8 +35,12 @@ func TestLargeOutputCompactsWithinAvailableContextAndContinues(t *testing.T) { partialSummary bool }{ {name: "completed summary", sourceStopReason: model.StopReasonMaxTokens}, - {name: "partial summary after answer cutoff", sourceStopReason: model.StopReasonMaxTokens, partialSummary: true}, - {name: "partial summary after completed answer", sourceStopReason: model.StopReasonEndTurn, partialSummary: true}, + {name: "partial summary after answer cutoff", + sourceStopReason: model.StopReasonMaxTokens, + partialSummary: true}, + {name: "partial summary after completed answer", + sourceStopReason: model.StopReasonEndTurn, + partialSummary: true}, } { t.Run(tc.name, func(t *testing.T) { ctx := context.Background() @@ -62,7 +66,11 @@ Finish any remaining work and provide a concise completion message.` } if tc.partialSummary { summaryResponse.StopReason = model.StopReasonMaxTokens - responses = append(responses, summaryResponse, completeProgressiveSummaryResponse("Finish the task.")) + responses = append( + responses, + summaryResponse, + completeProgressiveSummaryResponse("Finish the task."), + ) } responses = append(responses, summaryResponse, model.Response{ ID: "finished", StopReason: model.StopReasonEndTurn, @@ -70,23 +78,42 @@ Finish any remaining work and provide a concise completion message.` }) client := &sequenceKernelModel{ providerModelSlug: "large-output", - capabilities: model.Capabilities{ContextWindowTokens: 32_000, MaxOutputTokens: new(24_000)}, - responses: responses, + capabilities: model.Capabilities{ + ContextWindowTokens: 32_000, + MaxOutputTokens: new(24_000), + }, + responses: responses, + } + executor := AgentExecutor{ + Store: fixture.Store, + ModelResolver: liveTestModelResolver(fixture.Store, client), } - executor := AgentExecutor{Store: fixture.Store, ModelResolver: liveTestModelResolver(fixture.Store, client)} work := fixture.admitContentInputTurn(t, ctx, agentID, userID, "Complete the work and then confirm it is finished.", fixture.Now) for step := range 5 { require.NoError(t, executor.ExecuteModelWork(ctx, work)) fixture.releaseModelRuntimeLock(t, ctx, work) if step == 0 && tc.sourceStopReason == model.StopReasonEndTurn { - work = fixture.admitContentInputTurn(t, ctx, agentID, userID, "Continue the task.", fixture.Now.Add(time.Second)) + work = fixture.admitContentInputTurn( + t, + ctx, + agentID, + userID, + "Continue the task.", + fixture.Now.Add(time.Second), + ) continue } if pendingModelWork(t, ctx, fixture, agentID) == 0 { break } - claim := claimNextAgentWorkForKernelTest(t, ctx, fixture, agentID, executionstore.AgentWorkModel) + claim := claimNextAgentWorkForKernelTest( + t, + ctx, + fixture, + agentID, + executionstore.AgentWorkModel, + ) work = modelWorkExecutionFromClaimForKernelTest(claim, work.Now.Add(time.Second)) } wantRequests := 3 @@ -97,8 +124,11 @@ Finish any remaining work and provide a concise completion message.` require.Equal(t, 24_000, client.responded[0].Policy.MaxOutputTokens) largeSummary, continuation := client.responded[wantRequests-2], client.responded[wantRequests-1] require.True(t, isCompactionRequestBundle(largeSummary.Bundle)) - require.True(t, strings.Contains(string(largeSummary.ProviderRequest), strings.TrimSpace(largeText)), - "summary request must retain the complete large output") + require.True( + t, + strings.Contains(string(largeSummary.ProviderRequest), strings.TrimSpace(largeText)), + "summary request must retain the complete large output", + ) require.Positive(t, largeSummary.Policy.MaxOutputTokens) require.Less(t, largeSummary.Policy.MaxOutputTokens, 16_000) summaryInput := modelcontext.EstimatePreparedRequest(largeSummary.ProviderRequest, nil) @@ -110,9 +140,16 @@ Finish any remaining work and provide a concise completion message.` wantCutoffNotice := tc.sourceStopReason == model.StopReasonMaxTokens require.Equal(t, wantCutoffNotice, continuation.Bundle.ContextCheckpoint.EndsWithOutputLimit) checkpointText := modelcontext.ProjectedCheckpointContent(*continuation.Bundle.ContextCheckpoint) - require.Equal(t, wantCutoffNotice, strings.Contains(checkpointText, "[Automatic Omnara harness notice]")) - require.False(t, strings.Contains(string(continuation.ProviderRequest), strings.TrimSpace(largeText)), - "continuation should use the checkpoint") + require.Equal( + t, + wantCutoffNotice, + strings.Contains(checkpointText, "[Automatic Omnara harness notice]"), + ) + require.False( + t, + strings.Contains(string(continuation.ProviderRequest), strings.TrimSpace(largeText)), + "continuation should use the checkpoint", + ) require.Equal(t, 24_000, continuation.Policy.MaxOutputTokens) require.Zero(t, pendingModelWork(t, ctx, fixture, agentID)) var preserved, finished, failed, cutoffs int @@ -159,7 +196,10 @@ func TestOutputLimitContinuesAcrossClaimsUntilEndTurn(t *testing.T) { truncatedKernelResponse(), {ID: "done", StopReason: model.StopReasonEndTurn, Content: tc.content}, }} - executor := AgentExecutor{Store: fixture.Store, ModelResolver: liveTestModelResolver(fixture.Store, client)} + executor := AgentExecutor{ + Store: fixture.Store, + ModelResolver: liveTestModelResolver(fixture.Store, client), + } var openingInputs int require.NoError(t, fixture.Pool.QueryRow(ctx, `SELECT count(*) FROM agent_inputs WHERE agent_id=$1`, agentID).Scan(&openingInputs)) @@ -183,7 +223,13 @@ func TestOutputLimitContinuesAcrossClaimsUntilEndTurn(t *testing.T) { fixture.releaseModelRuntimeLock(t, ctx, work) if attempt < responseCount-1 { require.Equal(t, 1, pendingModelWork(t, ctx, fixture, agentID)) - claim := claimNextAgentWorkForKernelTest(t, ctx, fixture, agentID, executionstore.AgentWorkModel) + claim := claimNextAgentWorkForKernelTest( + t, + ctx, + fixture, + agentID, + executionstore.AgentWorkModel, + ) next := modelWorkExecutionFromClaimForKernelTest(claim, fixture.Now.Add(time.Second)) require.Equal(t, executionstore.ModelWorkContinue, next.Kind) require.Equal(t, work.TurnID, next.TurnID) @@ -221,7 +267,10 @@ func TestCancelOutputContinuationBeforeAndAfterRuntimeClaim(t *testing.T) { truncatedKernelResponse(), }, } - executor := AgentExecutor{Store: fixture.Store, ModelResolver: liveTestModelResolver(fixture.Store, client)} + executor := AgentExecutor{ + Store: fixture.Store, + ModelResolver: liveTestModelResolver(fixture.Store, client), + } require.NoError(t, executor.ExecuteModelWork(ctx, work)) fixture.releaseModelRuntimeLock(t, ctx, work) if afterClaim { @@ -275,7 +324,10 @@ func TestOutputContinuationIsConsumedBySuccessorTerminalFailure(t *testing.T) { truncatedKernelResponse(), }, } - executor := AgentExecutor{Store: fixture.Store, ModelResolver: liveTestModelResolver(fixture.Store, client)} + executor := AgentExecutor{ + Store: fixture.Store, + ModelResolver: liveTestModelResolver(fixture.Store, client), + } require.NoError(t, executor.ExecuteModelWork(ctx, work)) fixture.releaseModelRuntimeLock(t, ctx, work) client = &sequenceKernelModel{ @@ -332,7 +384,14 @@ func TestOutputContinuationSurvivesRetryAndCompaction(t *testing.T) { ), ModelRetryDelay: immediateKernelModelRetryDelay, } - work := fixture.admitContentInputTurn(t, ctx, agentID, userID, "establish earlier history", fixture.Now) + work := fixture.admitContentInputTurn( + t, + ctx, + agentID, + userID, + "establish earlier history", + fixture.Now, + ) require.NoError(t, executor.ExecuteModelWork(ctx, work)) fixture.releaseModelRuntimeLock(t, ctx, work) work = fixture.admitSteeringInputsTurn( @@ -419,7 +478,11 @@ func TestOutputContinuationSurvivesRetryAndCompaction(t *testing.T) { require.NotEmpty(t, request.Input) var checkpointText string require.NoError(t, json.Unmarshal(request.Input[0].Content, &checkpointText)) - require.Contains(t, checkpointText, "\n\n[Automatic Omnara harness notice]") + require.Contains( + t, + checkpointText, + "\n\n[Automatic Omnara harness notice]", + ) require.Greater(t, len(request.Input), 1) require.Contains(t, string(request.Input[1].Content), "complete the task") @@ -447,7 +510,10 @@ func TestNewInputSupersedesOutputContinuation(t *testing.T) { client := &sequenceKernelModel{providerModelSlug: "recovery-steering", responses: []model.Response{ truncatedKernelResponse(), {ID: "done", StopReason: model.StopReasonEndTurn}, }} - executor := AgentExecutor{Store: fixture.Store, ModelResolver: liveTestModelResolver(fixture.Store, client)} + executor := AgentExecutor{ + Store: fixture.Store, + ModelResolver: liveTestModelResolver(fixture.Store, client), + } work := fixture.admitContentInputTurn(t, ctx, agentID, userID, "complete the task", fixture.Now) require.NoError(t, executor.ExecuteModelWork(ctx, work)) fixture.releaseModelRuntimeLock(t, ctx, work) @@ -465,7 +531,11 @@ func TestNewInputSupersedesOutputContinuation(t *testing.T) { func pendingModelWork(t *testing.T, ctx context.Context, fixture kernelFixture, agentID uuid.UUID) int { t.Helper() var count int - require.NoError(t, fixture.Pool.QueryRow(ctx, - `SELECT count(*) FROM agent_next_model_work($1,$2)`, kernelTestProjectID, agentID).Scan(&count)) + require.NoError(t, fixture.Pool.QueryRow( + ctx, + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND h.agent_id=$2 AND h.logical_ready_at IS NOT NULL`, + kernelTestProjectID, + agentID, + ).Scan(&count)) return count } diff --git a/internal/harness/kernel/agent_executor_response.go b/internal/harness/kernel/agent_executor_response.go index b528152b0..5b885664b 100644 --- a/internal/harness/kernel/agent_executor_response.go +++ b/internal/harness/kernel/agent_executor_response.go @@ -22,6 +22,25 @@ func (e AgentExecutor) recordToolCallSourceEvent( specs []modelcontext.ToolSpec, streamedToolCallIDs map[string]uuid.UUID, ) (events.Event, error) { + bindings := toolCallBindings(envelope, specs, streamedToolCallIDs) + event, _, err := e.Store.Execution().RecordToolCallSourceAndCompleteContext( + ctx, + executionstore.RecordToolCallSourceAndCompleteContextInput{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + RuntimeLockID: input.RuntimeLockID, + ModelCallContextID: contextRow.ID, + ProviderRequestID: providerRequestID, + ProviderResponse: envelope, + ToolCallBindings: bindings, + }, + ) + return event, err +} + +func toolCallBindings( + envelope modelenvelope.ResponseEnvelope, specs []modelcontext.ToolSpec, streamedToolCallIDs map[string]uuid.UUID, +) []executionstore.ToolCallBindingInput { toolCalls := model.ToolCallsFromEnvelope(envelope) bindings := make([]executionstore.ToolCallBindingInput, 0, len(toolCalls)) toolsByName := toolSpecSet(specs) @@ -39,19 +58,7 @@ func (e AgentExecutor) recordToolCallSourceEvent( Type: toolType, }) } - event, _, err := e.Store.Execution().RecordToolCallSourceAndCompleteContext( - ctx, - executionstore.RecordToolCallSourceAndCompleteContextInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - RuntimeLockID: input.RuntimeLockID, - ModelCallContextID: contextRow.ID, - ProviderRequestID: providerRequestID, - ProviderResponse: envelope, - ToolCallBindings: bindings, - }, - ) - return event, err + return bindings } func invalidModelResponse(errorSource string, reason model.StopReason, calls []model.ToolCall) error { diff --git a/internal/harness/kernel/agent_executor_retry_integration_test.go b/internal/harness/kernel/agent_executor_retry_integration_test.go index e3565450c..1a6549edf 100644 --- a/internal/harness/kernel/agent_executor_retry_integration_test.go +++ b/internal/harness/kernel/agent_executor_retry_integration_test.go @@ -47,8 +47,10 @@ func TestAgentExecutorRetriesTransientProviderResponse(t *testing.T) { responses: []model.Response{{ ID: "resp-transient-retry", ProviderReportedCostUSD: "0.000003", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "done after retry"}}, - StopReason: model.StopReasonEndTurn, + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "done after retry"}, + }, + StopReason: model.StopReasonEndTurn, }}, errorResponses: []model.Response{ {ID: "resp-transient-error-1", ProviderReportedCostUSD: "0.000001"}, @@ -88,8 +90,14 @@ func TestAgentExecutorRetriesTransientProviderResponse(t *testing.T) { if err != nil { t.Fatalf("claim work for attempt %d: %v", attemptNumber, err) } - if !found || claim.Kind != executionstore.AgentWorkModel || claim.Model.ModelCallContextID == uuid.Nil { - t.Fatalf("attempt %d claim = %+v found=%v, want a retry continuation", attemptNumber, claim, found) + if !found || claim.Kind != executionstore.AgentWorkModel || + claim.Model.ModelCallContextID == uuid.Nil { + t.Fatalf( + "attempt %d claim = %+v found=%v, want a retry continuation", + attemptNumber, + claim, + found, + ) } predecessor, found, err := fixture.Store.Execution().GetModelCallContext( ctx, @@ -141,7 +149,12 @@ func TestAgentExecutorRetriesTransientProviderResponse(t *testing.T) { t.Fatalf("retry request %d changed bytes:\nfirst=%s\nretry=%s", index+1, firstRequest, got) } } - var failedContexts, contextsWithProviderMetadata, succeededContexts, contextsWithCost, contexts, operationFrontiers int + var failedContexts, + contextsWithProviderMetadata, + succeededContexts, + contextsWithCost, + contexts, + operationFrontiers int if err := fixture.Pool.QueryRow(ctx, ` SELECT count(*) FILTER ( WHERE context.state = 'failed' @@ -369,21 +382,18 @@ func TestManagedModelRetryReplayReturnsExistingAttemptAfterAdmissionCloses(t *te "create one retry attempt", fixture.Now.Add(time.Millisecond), ) - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, kernelTestProjectID, agentID) - if err != nil { - t.Fatalf("load managed retry replay agent: %v", err) - } - initial, err := fixture.Store.Execution().ClaimNormalModelCall( + + prepared1, err := fixture.Store.Execution().PrepareNormalModelCall( ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - RuntimeLockID: turn.RuntimeLockID, - OpeningInputIDs: turn.InputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: turn.OpeningEventSequence, + executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: agentID, + RuntimeLockID: turn.RuntimeLockID, + OpeningInputIDs: turn.InputIDs, }, ) + initial := prepared1.Claim + if err != nil { t.Fatalf("claim initial managed replay attempt: %v", err) } @@ -443,8 +453,10 @@ func TestAgentExecutorRetriesWithoutProviderReplayAfterReplayRejection(t *testin Message: "provider replay could not be decrypted", }}, responses: []model.Response{{ - ID: "resp-after-replay-rejection", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "continued canonically"}}, + ID: "resp-after-replay-rejection", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "continued canonically"}, + }, StopReason: model.StopReasonEndTurn, }}, } @@ -639,7 +651,10 @@ func TestAgentExecutorBoundsProviderEvidenceAttachedToError(t *testing.T) { modelClient := &sequenceKernelModel{ providerModelSlug: "error-evidence-model", errs: []error{model.ProviderError{ - Kind: model.ErrorKindTransient, Source: "test-provider", Code: "provider_unavailable", Message: "try later", + Kind: model.ErrorKindTransient, + Source: "test-provider", + Code: "provider_unavailable", + Message: "try later", }}, errorResponses: []model.Response{{ ID: strings.Repeat("r", model.MaxProviderIdentityBytes+1), @@ -680,7 +695,8 @@ WHERE context.project_id = $1 ); err != nil { t.Fatalf("load bounded provider error evidence: %v", err) } - if state != executionstore.ModelCallContextFailed || recoveryKind != executionstore.ModelCallRecoveryRetry || + if state != executionstore.ModelCallContextFailed || + recoveryKind != executionstore.ModelCallRecoveryRetry || responseID != "" { t.Fatalf( "provider error evidence = %q/%q response=%q", @@ -748,7 +764,8 @@ func TestAgentExecutorPreservesRetryAfterEvidenceWhenAttemptsAreExhausted(t *tes if err != nil { t.Fatalf("claim work for attempt %d: %v", attemptNumber, err) } - if !found || claim.Kind != executionstore.AgentWorkModel || claim.Model.ModelCallContextID == uuid.Nil { + if !found || claim.Kind != executionstore.AgentWorkModel || + claim.Model.ModelCallContextID == uuid.Nil { t.Fatalf("attempt %d claim = %+v found=%v, want retry continuation", attemptNumber, claim, found) } turn = modelWorkExecutionFromClaimForKernelTest(claim, currentNow) @@ -804,7 +821,14 @@ func TestAgentExecutorSteeringStartsFreshFrontierWithFreshRetryBudget(t *testing fixture := newKernelFixture(t, ctx) now := fixture.Now agentID, userID := fixture.createAgent(t, ctx, "openai/steering-retry-model", now) - turn := fixture.admitContentInputTurn(t, ctx, agentID, userID, "original request", now.Add(time.Millisecond)) + turn := fixture.admitContentInputTurn( + t, + ctx, + agentID, + userID, + "original request", + now.Add(time.Millisecond), + ) retryAfterSeconds := int64(3600) modelClient := &sequenceKernelModel{ providerModelSlug: "steering-retry-model", @@ -858,10 +882,12 @@ WHERE context.project_id = $1 steering, _, _, err := fixture.Store.Execution().CreateAgentContentInput( ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - Actor: kernelTestOmnaraActorParams(t, userID), - ContentBlocks: mustKernelJSON([]map[string]string{{"type": "text", "text": "steer to the new request"}}), + ProjectID: kernelTestProjectID, + AgentID: agentID, + Actor: kernelTestOmnaraActorParams(t, userID), + ContentBlocks: mustKernelJSON( + []map[string]string{{"type": "text", "text": "steer to the new request"}}, + ), DeliveryMode: executionstore.DeliveryModeSteering, IdempotencyKey: "steering-supersedes-retry", }, @@ -913,7 +939,8 @@ WHERE project_id = $1 AND agent_id = $2 AND id = $3 `, kernelTestProjectID, agentID, oldContextID).Scan(&oldState, &oldRecoveryKind); err != nil { t.Fatalf("load original retrying context: %v", err) } - if oldState != executionstore.ModelCallContextFailed || oldRecoveryKind != executionstore.ModelCallRecoveryRetry { + if oldState != executionstore.ModelCallContextFailed || + oldRecoveryKind != executionstore.ModelCallRecoveryRetry { t.Fatalf( "original retrying context = %q/%q, want immutable failed/retry history", oldState, @@ -968,8 +995,10 @@ func TestAgentExecutorConfigChangeRebuildsRetryingContextAtNewFrontier(t *testin }, }}, responses: []model.Response{{ - ID: "resp_after_config_change", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "continued with the new configuration"}}, + ID: "resp_after_config_change", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "continued with the new configuration"}, + }, StopReason: model.StopReasonEndTurn, }}, } @@ -1008,10 +1037,12 @@ WHERE context.project_id = $1 queuedInput, _, _, err := fixture.Store.Execution().CreateAgentContentInput( ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - Actor: kernelTestOmnaraActorParams(t, userID), - ContentBlocks: mustKernelJSON([]map[string]string{{"type": "text", "text": "wait for the active turn"}}), + ProjectID: kernelTestProjectID, + AgentID: agentID, + Actor: kernelTestOmnaraActorParams(t, userID), + ContentBlocks: mustKernelJSON( + []map[string]string{{"type": "text", "text": "wait for the active turn"}}, + ), DeliveryMode: executionstore.DeliveryModeQueued, IdempotencyKey: "queued-behind-config-change-retry", }, @@ -1021,7 +1052,8 @@ WHERE context.project_id = $1 } var configChangeStartedAt time.Time - if err := fixture.Pool.QueryRow(ctx, `SELECT statement_timestamp()`).Scan(&configChangeStartedAt); err != nil { + if err := fixture.Pool.QueryRow(ctx, + `SELECT statement_timestamp()`).Scan(&configChangeStartedAt); err != nil { t.Fatalf("read database time before config change: %v", err) } nextConfig := fixture.kernelAgentConfigInput( @@ -1041,7 +1073,8 @@ WHERE context.project_id = $1 t.Fatalf("change config during retry backoff: %v", err) } var configChangeCompletedAt time.Time - if err := fixture.Pool.QueryRow(ctx, `SELECT statement_timestamp()`).Scan(&configChangeCompletedAt); err != nil { + if err := fixture.Pool.QueryRow(ctx, + `SELECT statement_timestamp()`).Scan(&configChangeCompletedAt); err != nil { t.Fatalf("read database time after config change: %v", err) } var readyAt time.Time @@ -1053,7 +1086,8 @@ WHERE agent.project_id = $1 AND wake.agent_id = $2 `, kernelTestProjectID, agentID).Scan(&readyAt); err != nil { t.Fatalf("load wakeup after config change: %v", err) } - if readyAt.Before(configChangeStartedAt) || readyAt.After(configChangeCompletedAt) || !readyAt.Before(oldRetryAt) { + if readyAt.Before(configChangeStartedAt) || readyAt.After(configChangeCompletedAt) || + !readyAt.Before(oldRetryAt) { t.Fatalf( "config-change wakeup ready_at = %s, want within [%s, %s] before old retry %s", readyAt, @@ -1071,10 +1105,17 @@ WHERE agent.project_id = $1 AND wake.agent_id = $2 if err != nil { t.Fatalf("claim config-change continuation: %v", err) } - if !found || freshWork.Kind != executionstore.AgentWorkModel || freshWork.Model.ModelCallContextID != uuid.Nil || - freshWork.Model.TurnID != turn.TurnID || len(freshWork.Model.InputIDs) != 1 || + if !found || + freshWork.Kind != executionstore.AgentWorkModel || + freshWork.Model.ModelCallContextID != uuid.Nil || + freshWork.Model.TurnID != turn.TurnID || + len(freshWork.Model.InputIDs) != 1 || freshWork.Model.InputIDs[0] != turn.InputIDs[0] { - t.Fatalf("config-change work = %+v found=%v, want a fresh context for the active turn", freshWork, found) + t.Fatalf( + "config-change work = %+v found=%v, want a fresh context for the active turn", + freshWork, + found, + ) } oldRetry, err := fixture.Store.Execution().ClaimNextModelCallContext( @@ -1098,7 +1139,10 @@ WHERE agent.project_id = $1 AND wake.agent_id = $2 t.Fatalf("execute config-change continuation: %v", err) } if modelClient.preparedCount() != 2 { - t.Fatalf("prepared requests = %d, want failed old context and successful fresh context", modelClient.preparedCount()) + t.Fatalf( + "prepared requests = %d, want failed old context and successful fresh context", + modelClient.preparedCount(), + ) } var oldState, newState executionstore.ModelCallState @@ -1159,8 +1203,8 @@ WHERE project_id = $1 AND agent_id = $2 AND id = $3 var oldContinuableContexts int if err := fixture.Pool.QueryRow(ctx, ` SELECT count(*) -FROM agent_continuable_model_contexts($1, $2) -WHERE model_call_context_id = $3 +FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id +WHERE a.project_id=$1 AND h.agent_id=$2 AND h.normal_context_id=$3 AND h.turn_continuable `, kernelTestProjectID, agentID, oldContextID).Scan(&oldContinuableContexts); err != nil { t.Fatalf("check stale context continuation eligibility: %v", err) } @@ -1229,13 +1273,17 @@ func TestAgentExecutorQueuedInputWaitsForRetryingTurn(t *testing.T) { }}, responses: []model.Response{ { - ID: "resp_after_retry", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "finished the original request"}}, + ID: "resp_after_retry", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "finished the original request"}, + }, StopReason: model.StopReasonEndTurn, }, { - ID: "resp_after_queued_input", - Content: []model.ResponsePart{{Type: model.ResponsePartTypeText, Text: "handled the queued request"}}, + ID: "resp_after_queued_input", + Content: []model.ResponsePart{ + {Type: model.ResponsePartTypeText, Text: "handled the queued request"}, + }, StopReason: model.StopReasonEndTurn, }, }, @@ -1316,7 +1364,8 @@ WHERE project_id = $1 AND agent_id = $2 AND id = $3 if err != nil { t.Fatalf("claim retry at deadline: %v", err) } - if !found || retryWork.Kind != executionstore.AgentWorkModel || retryWork.Model.ModelCallContextID != oldContextID || + if !found || retryWork.Kind != executionstore.AgentWorkModel || + retryWork.Model.ModelCallContextID != oldContextID || retryWork.Model.TurnID != turn.TurnID { t.Fatalf("retry work = %+v found=%v, want original turn/context", retryWork, found) } @@ -1366,7 +1415,10 @@ WHERE agent.project_id = $1 AND wake.agent_id = $2 t.Fatalf("execute queued input: %v", err) } if modelClient.preparedCount() != 3 { - t.Fatalf("prepared requests = %d, want initial attempt, retry, and queued turn", modelClient.preparedCount()) + t.Fatalf( + "prepared requests = %d, want initial attempt, retry, and queued turn", + modelClient.preparedCount(), + ) } var oldState, queuedContextState executionstore.ModelCallState var oldContextCount, succeededRetryContexts int @@ -1406,9 +1458,12 @@ WHERE context.project_id = $1 ); err != nil { t.Fatalf("load queued turn model context: %v", err) } - if oldState != executionstore.ModelCallContextFailed || oldContextCount != 2 || succeededRetryContexts != 1 || + if oldState != executionstore.ModelCallContextFailed || + oldContextCount != 2 || + succeededRetryContexts != 1 || newContextID == oldContextID || - queuedContextState != executionstore.ModelCallContextSucceeded || newAttemptNumber != 1 { + queuedContextState != executionstore.ModelCallContextSucceeded || + newAttemptNumber != 1 { t.Fatalf( "contexts old=%s/%s lineage_rows=%d succeeded_retries=%d queued=%s/%s attempt=%d, want immutable failed predecessor plus succeeded retry, then a fresh succeeded context", oldContextID, diff --git a/internal/harness/kernel/agent_executor_retry_runtime_recovery_integration_test.go b/internal/harness/kernel/agent_executor_retry_runtime_recovery_integration_test.go index 12b01d667..df2c6fb63 100644 --- a/internal/harness/kernel/agent_executor_retry_runtime_recovery_integration_test.go +++ b/internal/harness/kernel/agent_executor_retry_runtime_recovery_integration_test.go @@ -58,7 +58,8 @@ WHERE context.project_id = $1 t.Fatalf("load malformed response attempt: %v", err) } if state != executionstore.ModelCallContextFailed || recoveryKind != executionstore.ModelCallRecoveryRetry || - errorCode != "malformed_success_response" || strings.Contains(errorMessage, "\x00") || + errorCode != "malformed_success_response" || + strings.Contains(errorMessage, "\x00") || !kernelModelCallOutcomeAmbiguous(t, errorDetails) { t.Fatalf( "malformed response attempt = %q/%q/%q message=%q", @@ -96,24 +97,20 @@ func TestAgentExecutorRetriesInterruptedModelContext(t *testing.T) { "recover interrupted model call", fixture.Now.Add(time.Second), ) - watermark, err := fixture.Store.Execution().MaxEventSequence(ctx, kernelTestProjectID, agentID) + _, err := fixture.Store.Execution().MaxEventSequence(ctx, kernelTestProjectID, agentID) if err != nil { t.Fatalf("max event sequence: %v", err) } - snapshot, err := fixture.Store.Execution().CaptureAgentConfigForEventWatermark( - ctx, kernelTestProjectID, agentID, watermark, - ) - if err != nil { - t.Fatalf("capture config snapshot: %v", err) - } - modelClaim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - RuntimeLockID: turn.RuntimeLockID, - OpeningInputIDs: turn.InputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: watermark, - }) + + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: kernelTestProjectID, + AgentID: agentID, + RuntimeLockID: turn.RuntimeLockID, + OpeningInputIDs: turn.InputIDs, + }) + modelClaim := prepared1.Claim + if err != nil { t.Fatalf("claim initial model context: %v", err) } @@ -184,7 +181,10 @@ WHERE project_id = $1 AND operation_kind = 'normal' ORDER BY attempt_number DESC LIMIT 1 -`, kernelTestProjectID, agentID, modelClaim.Context.InputEventSequence).Scan(&recoveredContextID); err != nil { +`, + kernelTestProjectID, + agentID, + modelClaim.Context.InputEventSequence).Scan(&recoveredContextID); err != nil { t.Fatalf("load recovered context id: %v", err) } recovered, found, err := fixture.Store.Execution().GetModelCallContext( @@ -212,7 +212,11 @@ LIMIT 1 AND agent_id = $2 AND input_event_sequence = $3 AND operation_kind = 'normal' - `, kernelTestProjectID, agentID, modelClaim.Context.InputEventSequence).Scan(&interruptedAttempts, &succeededAttempts); err != nil { + `, + kernelTestProjectID, + agentID, + modelClaim.Context.InputEventSequence).Scan(&interruptedAttempts, + &succeededAttempts); err != nil { t.Fatalf("count interrupted context attempts: %v", err) } if interruptedAttempts != 1 || succeededAttempts != 1 { diff --git a/internal/harness/kernel/agent_executor_test_helpers_test.go b/internal/harness/kernel/agent_executor_test_helpers_test.go index 6f8d82d23..519b64323 100644 --- a/internal/harness/kernel/agent_executor_test_helpers_test.go +++ b/internal/harness/kernel/agent_executor_test_helpers_test.go @@ -222,7 +222,8 @@ func (f kernelFixture) createNamedAgentWithModelOptions( t.Helper() providerConfigBaseName, configuredModelName, ok := strings.Cut(modelSelection, "/") if !ok { - t.Fatalf("kernel test model selection %q must be provider-config-base/configured-model-name", modelSelection) + t.Fatalf("kernel test model selection %q must be provider-config-base/configured-model-name", + modelSelection) } providerConfigName := providerConfigBaseName + "-prod" agentProfileIdempotencyKey := "kernel-agent-" + providerConfigName + "-" + configuredModelName + "-" + now.Format( @@ -331,41 +332,48 @@ func (f kernelFixture) compileAgentYAMLResolvedWithModelOptions( if err != nil { t.Fatalf("parse agent config source: %v", err) } - configuredModel := f.ensureModelSelection(t, ctx, source.Model.ProviderConfig, source.Model.Name, modelOptions) - compiled, err := agentconfig.Compile(agentconfig.SourceFormatYAML, []byte(sourceYAML), agentconfig.CompileOptions{ - ResolveModelSelection: func( - providerConfigName string, - configuredModelName string, - ) (agentconfig.ResolvedModelSelection, error) { - return resolvedKernelAgentConfigModel(configuredModel), nil - }, - ResolveMachinePoolName: func(machinePoolName string) (uuid.UUID, error) { - machinePoolID, err := f.Store.Execution().ResolveAgentConfigMachinePoolName( - ctx, - kernelTestOrgID, - kernelTestProjectID, - machinePoolName, - ) - if err != nil { - return uuid.Nil, err - } - return machinePoolID, nil - }, - ResolveSkillID: func(skillID string) (agentconfig.SkillResolution, error) { - records, _, err := f.Store.Skills().GetSkillsByIDsForCompile(ctx, skillstore.GetSkillsByIDsInput{ - OrgID: kernelTestOrgID, - ProjectID: kernelTestProjectID, - IDs: []string{skillID}, - }) - if err != nil { - return agentconfig.SkillResolution{}, err - } - if len(records) != 1 { - return agentconfig.SkillResolution{}, storeerr.ErrNotFound - } - return agentconfig.SkillResolution{ID: records[0].ID, Name: records[0].Name}, nil - }, - }) + configuredModel := f.ensureModelSelection(t, + ctx, + source.Model.ProviderConfig, + source.Model.Name, + modelOptions) + compiled, + err := agentconfig.Compile(agentconfig.SourceFormatYAML, + []byte(sourceYAML), + agentconfig.CompileOptions{ + ResolveModelSelection: func( + providerConfigName string, + configuredModelName string, + ) (agentconfig.ResolvedModelSelection, error) { + return resolvedKernelAgentConfigModel(configuredModel), nil + }, + ResolveMachinePoolName: func(machinePoolName string) (uuid.UUID, error) { + machinePoolID, err := f.Store.Execution().ResolveAgentConfigMachinePoolName( + ctx, + kernelTestOrgID, + kernelTestProjectID, + machinePoolName, + ) + if err != nil { + return uuid.Nil, err + } + return machinePoolID, nil + }, + ResolveSkillID: func(skillID string) (agentconfig.SkillResolution, error) { + records, _, err := f.Store.Skills().GetSkillsByIDsForCompile(ctx, skillstore.GetSkillsByIDsInput{ + OrgID: kernelTestOrgID, + ProjectID: kernelTestProjectID, + IDs: []string{skillID}, + }) + if err != nil { + return agentconfig.SkillResolution{}, err + } + if len(records) != 1 { + return agentconfig.SkillResolution{}, storeerr.ErrNotFound + } + return agentconfig.SkillResolution{ID: records[0].ID, Name: records[0].Name}, nil + }, + }) if err != nil { t.Fatalf("compile resolved agent config: %v", err) } @@ -390,7 +398,9 @@ func (f kernelFixture) ensureModelSelection( ) modelstore.ConfiguredModelRecord { t.Helper() provider := storagefixture.EnsureModelProvider(t, ctx, f.Store.Models(), f.Store.Secrets(), - storagefixture.ModelProviderInput{OrgID: kernelTestOrgID, UserID: kernelTestUserID, Name: providerConfigName}) + storagefixture.ModelProviderInput{OrgID: kernelTestOrgID, + UserID: kernelTestUserID, + Name: providerConfigName}) if provider.ManagementKind == management.Cluster { configuredModel, err := f.Store.Models().GetConfiguredModelByName( ctx, kernelTestOrgID, provider.ID, configuredModelName, @@ -534,7 +544,11 @@ func currentProjectModelGrantIDForKernelConfiguredModelID( configuredModelID uuid.UUID, ) uuid.UUID { t.Helper() - grant, err := store.Models().GetActiveProjectModelGrantForConfiguredModel(ctx, orgID, projectID, configuredModelID) + grant, + err := store.Models().GetActiveProjectModelGrantForConfiguredModel(ctx, + orgID, + projectID, + configuredModelID) if err != nil { t.Fatalf("load project model grant for configured model %s: %v", configuredModelID, err) } @@ -594,13 +608,17 @@ func (f kernelFixture) admitContentInputTurn( now time.Time, ) ModelWorkExecution { t.Helper() - input, _, _, err := f.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: kernelTestProjectID, - AgentID: agentID, - Actor: kernelTestOmnaraActorParams(t, userID), - ContentBlocks: mustKernelJSON([]map[string]string{{"type": "text", "text": text}}), - IdempotencyKey: "kernel-input-" + agentID.String() + "-" + text, - }) + input, + _, + _, + err := f.Store.Execution().CreateAgentContentInput(ctx, + executionstore.CreateAgentContentInputInput{ + ProjectID: kernelTestProjectID, + AgentID: agentID, + Actor: kernelTestOmnaraActorParams(t, userID), + ContentBlocks: mustKernelJSON([]map[string]string{{"type": "text", "text": text}}), + IdempotencyKey: "kernel-input-" + agentID.String() + "-" + text, + }) if err != nil { t.Fatalf("create agent input: %v", err) } @@ -629,7 +647,8 @@ func (f kernelFixture) admitContentInputTurn( } } -func modelWorkExecutionFromClaimForKernelTest(claim executionstore.ClaimedAgentWork, now time.Time) ModelWorkExecution { +func modelWorkExecutionFromClaimForKernelTest(claim executionstore.ClaimedAgentWork, + now time.Time) ModelWorkExecution { return ModelWorkExecution{ Kind: claim.Model.Kind, OrgID: claim.OrgID, @@ -834,7 +853,9 @@ func continueTurnOnNewLeaseForKernelTest( prior.AgentID, ).Scan(&checkpoints) _ = fixture.Pool.QueryRow( - ctx, `SELECT count(*) FROM agent_unconsumed_context_checkpoint_frontiers($1, $2)`, prior.ProjectID, + ctx, + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND h.agent_id=$2 AND h.pending_checkpoint_id IS NOT NULL`, + prior.ProjectID, prior.AgentID, ).Scan(&checkpointFrontiers) _ = fixture.Pool.QueryRow( @@ -971,7 +992,9 @@ func (m *sequenceKernelModel) Capabilities() model.Capabilities { return capabilities } -func (m *sequenceKernelModel) Prepare(_ context.Context, input model.PrepareInput) (model.PreparedRequest, error) { +func (m *sequenceKernelModel) Prepare(_ context.Context, + input model.PrepareInput) (model.PreparedRequest, + error) { if m.prepareErr != nil { return model.PreparedRequest{}, m.prepareErr } @@ -1250,7 +1273,10 @@ func (c *fakeKernelMCPClient) Initialize( }, nil } -func (c *fakeKernelMCPClient) Notify(_ context.Context, conn mcp.Conn, method string, _ json.RawMessage) error { +func (c *fakeKernelMCPClient) Notify(_ context.Context, + conn mcp.Conn, + method string, + _ json.RawMessage) error { c.mu.Lock() c.notifyCount++ agentID := c.expectedAgentIDLocked() @@ -1264,7 +1290,12 @@ func (c *fakeKernelMCPClient) Notify(_ context.Context, conn mcp.Conn, method st return nil } -func (c *fakeKernelMCPClient) Call(context.Context, mcp.Conn, string, json.RawMessage, int64) (json.RawMessage, error) { +func (c *fakeKernelMCPClient) Call(context.Context, + mcp.Conn, + string, + json.RawMessage, + int64) (json.RawMessage, + error) { return nil, errors.New("unexpected generic mcp call") } diff --git a/internal/harness/kernel/agent_executor_tool_lifecycle_integration_test.go b/internal/harness/kernel/agent_executor_tool_lifecycle_integration_test.go index 0e1256f2c..12a104831 100644 --- a/internal/harness/kernel/agent_executor_tool_lifecycle_integration_test.go +++ b/internal/harness/kernel/agent_executor_tool_lifecycle_integration_test.go @@ -130,6 +130,12 @@ skills: } } +func (e *requeueingToolWorkExecutor) DispatchRecordedToolCall( + ctx context.Context, turn tools.Turn, record executionstore.ToolCallRecord, +) (tools.Result, error) { + return e.Dispatch(ctx, turn, modelToolCallFromRecord(record)) +} + func (e *requeueingToolWorkExecutor) Dispatch( ctx context.Context, turn tools.Turn, diff --git a/internal/harness/kernel/agent_executor_tool_work.go b/internal/harness/kernel/agent_executor_tool_work.go index 2b339415c..ae921fb3f 100644 --- a/internal/harness/kernel/agent_executor_tool_work.go +++ b/internal/harness/kernel/agent_executor_tool_work.go @@ -8,8 +8,10 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/harness/tools" "github.com/omnara-ai/omnara/internal/model" + "github.com/omnara-ai/omnara/internal/modelcontext" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/toolcatalog" ) func (e AgentExecutor) ExecuteToolWork(ctx context.Context, input ToolWorkExecution) error { @@ -22,6 +24,13 @@ type toolWorkExecutor interface { FailRunnableToolCall(context.Context, tools.Turn, model.ToolCall, error) error } +type recordedToolWorkExecutor interface { + PrepareRecordedToolCallPermission( + context.Context, tools.Turn, executionstore.ToolCallRecord, + ) (executionstore.ToolCallRecord, error) + DispatchRecordedToolCall(context.Context, tools.Turn, executionstore.ToolCallRecord) (tools.Result, error) +} + func (e AgentExecutor) executeToolWork( ctx context.Context, input ToolWorkExecution, @@ -39,56 +48,78 @@ func (e AgentExecutor) executeToolWork( if input.Now.IsZero() { input.Now = e.now() } - contextRecord, found, err := e.Store.Execution().GetModelCallContext( - ctx, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) - if err != nil { - return fmt.Errorf("load tool work model context: %w", err) - } - if !found { - return fmt.Errorf("tool work model context %s not found", input.ModelCallContextID) - } - output, found, err := e.Store.Execution().GetModelOutputForContext( - ctx, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) - if err != nil { - return fmt.Errorf("load tool work model output: %w", err) - } - if contextRecord.State != executionstore.ModelCallContextSucceeded || - !found || - output.ID != input.ModelOutputID { - return fmt.Errorf( - "tool work does not match accepted model output for context %s: %w", + var contextRecord executionstore.ModelCallContextRecord + var specs []modelcontext.ToolSpec + var first *executionstore.ToolCallRecord + var err error + if input.Prepared != nil { + contextRecord = input.Prepared.Context + if contextRecord.ID != input.ModelCallContextID || contextRecord.AgentID != input.AgentID || + contextRecord.ProjectID != input.ProjectID || contextRecord.AgentConfigID != input.Prepared.Config.ID { + return storeerr.ErrStateTransitionConflict + } + specs, err = e.toolRuntimeFromConfig(ctx, input.ProjectID, input.AgentID, input.Prepared.Config, input.Now) + first = input.Prepared.FirstCall + } else { + var found bool + contextRecord, found, err = e.Store.Execution().GetModelCallContext( + ctx, + input.ProjectID, + input.AgentID, input.ModelCallContextID, - storeerr.ErrStateTransitionConflict, + ) + if err != nil { + return fmt.Errorf("load tool work model context: %w", err) + } + if !found { + return fmt.Errorf("tool work model context %s not found", input.ModelCallContextID) + } + output, found, outputErr := e.Store.Execution().GetModelOutputForContext( + ctx, + input.ProjectID, + input.AgentID, + input.ModelCallContextID, + ) + if outputErr != nil { + return fmt.Errorf("load tool work model output: %w", outputErr) + } + if contextRecord.State != executionstore.ModelCallContextSucceeded || + !found || + output.ID != input.ModelOutputID { + return fmt.Errorf( + "tool work does not match accepted model output for context %s: %w", + input.ModelCallContextID, + storeerr.ErrStateTransitionConflict, + ) + } + specs, err = e.modelContextToolRuntime( + ctx, + input.ProjectID, + input.AgentID, + contextRecord, + input.Now, ) } - specs, err := e.modelContextToolRuntime( - ctx, - input.ProjectID, - input.AgentID, - contextRecord, - input.Now, - ) if err != nil { return fmt.Errorf("load tool work runtime contract: %w", err) } turn := toolWorkTurn(input, contextRecord.OrgID, specs) deferredToolCallIDs := make([]uuid.UUID, 0) for { - record, found, err := e.Store.Execution().NextRunnableToolCall( - ctx, - input.ProjectID, - input.AgentID, - input.ModelOutputID, - deferredToolCallIDs, - ) + var record executionstore.ToolCallRecord + var found bool + if first != nil { + record, found = *first, true + first = nil + } else { + record, found, err = e.Store.Execution().NextRunnableToolCall( + ctx, + input.ProjectID, + input.AgentID, + input.ModelOutputID, + deferredToolCallIDs, + ) + } if err != nil { return err } @@ -105,6 +136,43 @@ func (e AgentExecutor) executeToolWork( ) } call := modelToolCallFromRecord(record) + if recorded, ok := executor.(recordedToolWorkExecutor); ok { + if record.State == executionstore.ToolCallStateAwaitingAuthorization { + prepared, err := recorded.PrepareRecordedToolCallPermission(ctx, turn, record) + if err != nil { + return fmt.Errorf("prepare permission for tool %q: %w", call.ID, err) + } + if prepared.ID == uuid.Nil { + prepared, err = e.Store.Execution().GetToolCall(ctx, input.ProjectID, input.AgentID, record.ID) + if err != nil { + return err + } + } + if prepared.State == record.State { + if err := executor.FailRunnableToolCall(ctx, turn, call, + fmt.Errorf("tool coordinator returned without moving tool call out of %s", record.State)); err != nil { + return err + } + continue + } + record = prepared + } + if record.State != executionstore.ToolCallStateReady || !toolcatalog.IsPlatformManagedToolType(record.Type) { + continue + } + result, err := recorded.DispatchRecordedToolCall(ctx, turn, record) + if err != nil { + return fmt.Errorf("dispatch tool %q: %w", call.ID, err) + } + switch result.Disposition { + case tools.DispatchDeferred: + deferredToolCallIDs = append(deferredToolCallIDs, record.ID) + case tools.DispatchCompleted: + default: + return fmt.Errorf("dispatch tool %q returned invalid disposition %d", call.ID, result.Disposition) + } + continue + } switch record.State { case executionstore.ToolCallStateAwaitingAuthorization: if err := executor.PrepareToolCallPermission(ctx, turn, call); err != nil { diff --git a/internal/harness/kernel/agent_executor_web_tools_integration_test.go b/internal/harness/kernel/agent_executor_web_tools_integration_test.go index 627d428dd..15e17943a 100644 --- a/internal/harness/kernel/agent_executor_web_tools_integration_test.go +++ b/internal/harness/kernel/agent_executor_web_tools_integration_test.go @@ -75,7 +75,12 @@ func TestAgentExecutorWebToolsAsyncLifecycle(t *testing.T) { }, Now: func() time.Time { return fixture.Now.Add(2 * time.Second) }, } - turn := fixture.admitContentInputTurn(t, ctx, agentID, userID, "research go releases", fixture.Now.Add(time.Second)) + turn := fixture.admitContentInputTurn(t, + ctx, + agentID, + userID, + "research go releases", + fixture.Now.Add(time.Second)) if err := executor.ExecuteModelWork(ctx, turn); err != nil { t.Fatalf("persist web tool output: %v", err) } @@ -218,7 +223,9 @@ func (g *gatedSearchProvider) Search( } return webaccess.SearchResponse{ Provider: "gated", - Results: []webaccess.SearchResult{{URL: "https://example.org/gated", Title: "Gated", Snippet: "gated result"}}, + Results: []webaccess.SearchResult{{URL: "https://example.org/gated", + Title: "Gated", + Snippet: "gated result"}}, }, nil } @@ -287,7 +294,10 @@ WHERE project_id = $1 var runnable int if err := fixture.Pool.QueryRow(ctx, ` SELECT count(*) -FROM agent_tool_work_frontiers($1, $2) +FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id +WHERE a.project_id=$1 AND h.agent_id=$2 AND EXISTS(SELECT 1 FROM tool_calls t + WHERE t.agent_id=h.agent_id AND t.model_output_id=h.pending_tool_output_id + AND (t.state='awaiting_authorization' OR t.state='ready' AND t.type IN ('built_in','mcp'))) `, kernelTestProjectID, agentID).Scan(&runnable); err != nil { t.Fatalf("count runnable tool work: %v", err) } diff --git a/internal/harness/kernel/types.go b/internal/harness/kernel/types.go index e9e3d9881..a2f1a44e0 100644 --- a/internal/harness/kernel/types.go +++ b/internal/harness/kernel/types.go @@ -8,6 +8,7 @@ import ( ) type ModelWorkExecution struct { + Prepared *executionstore.PreparedNormalModelCall Kind executionstore.ModelWorkKind OrgID uuid.UUID ProjectID uuid.UUID @@ -23,6 +24,7 @@ type ModelWorkExecution struct { } type ToolWorkExecution struct { + Prepared *executionstore.PreparedToolWork ProjectID uuid.UUID AgentID uuid.UUID TurnID uuid.UUID @@ -32,3 +34,8 @@ type ToolWorkExecution struct { RuntimeLockID uuid.UUID Now time.Time } + +type ModelWorkAdvanceOptions struct { + Deadline time.Time + AllowModelWork bool +} diff --git a/internal/harness/tools/approvals.go b/internal/harness/tools/approvals.go index 6200ca973..a8f1aef8f 100644 --- a/internal/harness/tools/approvals.go +++ b/internal/harness/tools/approvals.go @@ -20,20 +20,29 @@ func (e Executor) PrepareToolCallPermission( turn Turn, call model.ToolCall, ) error { - if _, ok, err := e.completedToolResult(ctx, turn, call); err != nil || ok { - return err - } proposal, err := e.recordedToolCall(ctx, turn, call) if err != nil { return err } + _, err = e.PrepareRecordedToolCallPermission(ctx, turn, proposal) + return err +} + +func (e Executor) PrepareRecordedToolCallPermission( + ctx context.Context, turn Turn, proposal executionstore.ToolCallRecord, +) (executionstore.ToolCallRecord, error) { + if proposal.ProjectID != turn.ProjectID || proposal.AgentID != turn.AgentID || + proposal.ModelCallContextID != turn.ModelCallContextID { + return executionstore.ToolCallRecord{}, storeerr.ErrIdempotencyConflict + } + call := model.ToolCall{ID: proposal.ProviderCallID, Name: proposal.Name, Input: proposal.Input} if proposal.State != executionstore.ToolCallStateAwaitingAuthorization && proposal.State != executionstore.ToolCallStateAwaitingPermission { - return nil + return proposal, nil } spec, configured := turn.Tools[call.Name] if !configured { - return e.completeInvalidToolCall( + return executionstore.ToolCallRecord{}, e.completeInvalidToolCall( ctx, turn, proposal.ID, @@ -45,7 +54,7 @@ func (e Executor) PrepareToolCallPermission( toolType := spec.Type implementation, implemented, err := toolImplementationFor(call.Name) if err != nil { - return err + return executionstore.ToolCallRecord{}, err } modeHandler, modeDescriptor, supported, err := permissionModeForTool( toolType, @@ -54,7 +63,7 @@ func (e Executor) PrepareToolCallPermission( implemented, ) if err != nil { - return err + return executionstore.ToolCallRecord{}, err } var inputErr error if implemented { @@ -76,10 +85,10 @@ func (e Executor) PrepareToolCallPermission( errorCode = "malformed" } cause := firstError(inputErr, unsupportedErr) - return e.completeInvalidToolCall(ctx, turn, proposal.ID, errorCode, cause) + return executionstore.ToolCallRecord{}, e.completeInvalidToolCall(ctx, turn, proposal.ID, errorCode, cause) } if !supported { - return fmt.Errorf( + return executionstore.ToolCallRecord{}, fmt.Errorf( "tool %q does not support permission mode %q", call.Name, selection.Mode, @@ -90,7 +99,7 @@ func (e Executor) PrepareToolCallPermission( []toolpermission.ModeDescriptor{modeDescriptor}, ) if err != nil { - return fmt.Errorf("tool %q runtime permission: %w", call.Name, err) + return executionstore.ToolCallRecord{}, fmt.Errorf("tool %q runtime permission: %w", call.Name, err) } if existing, found, err := e.Store.Execution().GetAgentInteractionByToolCallKind( ctx, @@ -99,9 +108,9 @@ func (e Executor) PrepareToolCallPermission( proposal.ID, executionstore.AgentInteractionKindPermission, ); err != nil { - return err + return executionstore.ToolCallRecord{}, err } else if found { - return e.evaluateExistingPermissionInteraction( + return executionstore.ToolCallRecord{}, e.evaluateExistingPermissionInteraction( ctx, turn, call, @@ -111,7 +120,7 @@ func (e Executor) PrepareToolCallPermission( ) } if proposal.State == executionstore.ToolCallStateAwaitingPermission { - return storeerr.ErrIdempotencyConflict + return executionstore.ToolCallRecord{}, storeerr.ErrIdempotencyConflict } modeResult, err := modeHandler( ctx, @@ -125,7 +134,7 @@ func (e Executor) PrepareToolCallPermission( ) var preparationErr *toolCallPreparationError if errors.As(err, &preparationErr) { - return e.completeToolCallPreparationFailure( + return executionstore.ToolCallRecord{}, e.completeToolCallPreparationFailure( ctx, turn, proposal.ID, @@ -133,21 +142,15 @@ func (e Executor) PrepareToolCallPermission( ) } if err != nil { - return err + return executionstore.ToolCallRecord{}, err } switch modeResult.kind { case permissionModeAllow: - if _, err := e.Store.Execution().MarkToolCallReady(ctx, executionstore.MarkToolCallReadyInput{ - ProjectID: turn.ProjectID, - AgentID: turn.AgentID, - ID: proposal.ID, - RuntimeLockID: turn.RuntimeLockID, - }); err != nil { - return err - } - return nil + return e.Store.Execution().MarkToolCallReady(ctx, executionstore.MarkToolCallReadyInput{ + ProjectID: turn.ProjectID, AgentID: turn.AgentID, ID: proposal.ID, RuntimeLockID: turn.RuntimeLockID, + }) case permissionModeDeny: - return e.completeDeniedToolCall( + return executionstore.ToolCallRecord{}, e.completeDeniedToolCall( ctx, turn, proposal.ID, @@ -155,7 +158,7 @@ func (e Executor) PrepareToolCallPermission( ) case permissionModeAsk: default: - return fmt.Errorf("tool %q permission mode returned an invalid outcome", call.Name) + return executionstore.ToolCallRecord{}, fmt.Errorf("tool %q permission mode returned an invalid outcome", call.Name) } interaction, err := e.Store.Execution().CreatePermissionInteraction( ctx, @@ -168,10 +171,10 @@ func (e Executor) PrepareToolCallPermission( }, ) if err != nil { - return err + return executionstore.ToolCallRecord{}, err } e.enqueueIntegrationPromptCopy(turn, interaction) - return nil + return executionstore.ToolCallRecord{}, nil } func (e Executor) evaluateExistingPermissionInteraction( diff --git a/internal/harness/tools/dispatch.go b/internal/harness/tools/dispatch.go index 988f6bf16..c99bc2331 100644 --- a/internal/harness/tools/dispatch.go +++ b/internal/harness/tools/dispatch.go @@ -204,7 +204,10 @@ func (p toolPhasePipeline) advanceAfterTransaction( execution.CommandResult, ) } - return toolDispatchCompleted{}, false, nil + if execution.Completed == nil { + return nil, false, errors.New("completed tool transaction did not return its result") + } + return toolDispatchCompleted{record: *execution.Completed}, false, nil case executionstore.ToolCallDispositionRunning: default: return nil, false, fmt.Errorf( @@ -651,7 +654,7 @@ type toolDispatchResult interface { } type toolDispatchAwaiting struct{} -type toolDispatchCompleted struct{} +type toolDispatchCompleted struct{ record executionstore.ToolCallRecord } type toolDispatchFailed struct { content toolResultContent diff --git a/internal/harness/tools/dispatch_transaction_integration_test.go b/internal/harness/tools/dispatch_transaction_integration_test.go index b64c58e32..ac5dd9378 100644 --- a/internal/harness/tools/dispatch_transaction_integration_test.go +++ b/internal/harness/tools/dispatch_transaction_integration_test.go @@ -86,17 +86,25 @@ func TestListMachinesDispatchRejectsInvalidCursor(t *testing.T) { if err != nil { t.Fatal(err) } - if result.Disposition != DispatchCompleted || - !strings.Contains(string(result.ContentParts), `"error_code":"malformed"`) { + if result.Disposition != DispatchCompleted { t.Fatalf("invalid cursor result = %+v, want completed malformed-input result", result) } + var parts []struct { + Type string + Value map[string]any + } + require.NoError(t, json.Unmarshal(result.ContentParts, &parts)) + require.Len(t, parts, 1) + require.Equal(t, "structured_data", parts[0].Type) + require.Equal(t, "malformed", parts[0].Value["error_code"]) record, err := fixture.Store.Execution().GetToolCall( ctx, toolsTestProjectID, fixture.Agent.ID, fixture.toolCallID(t, ctx, call.ID), ) if err != nil { t.Fatal(err) } - if record.State != executionstore.ToolCallStateCompleted || record.Outcome != executionstore.ToolResultOutcomeFailed { + if record.State != executionstore.ToolCallStateCompleted || + record.Outcome != executionstore.ToolResultOutcomeFailed { t.Fatalf("invalid cursor tool call = %+v, want completed failure", record) } } @@ -171,7 +179,14 @@ func TestSpawnAgentDispatchReturnsUnsupportedModelToolFailure(t *testing.T) { result, err := (Executor{Store: fixture.Store}).Dispatch(ctx, turn, call) require.NoError(t, err) require.Equal(t, DispatchCompleted, result.Disposition) - require.Contains(t, string(result.ContentParts), `"error_code":"spawn_agent_failed"`) + var parts []struct { + Type string + Value map[string]any + } + require.NoError(t, json.Unmarshal(result.ContentParts, &parts)) + require.Len(t, parts, 1) + require.Equal(t, "structured_data", parts[0].Type) + require.Equal(t, "spawn_agent_failed", parts[0].Value["error_code"]) require.Contains(t, string(result.ContentParts), "does not support tools") record, err := fixture.Store.Execution().GetToolCall( ctx, toolsTestProjectID, fixture.Agent.ID, fixture.toolCallID(t, ctx, call.ID), @@ -1285,7 +1300,12 @@ func TestFileToolsApprovalDispatch(t *testing.T) { fixture := newIntegrationToolFixtureWithOptions( t, ctx, "file-approval-"+name, toolFixtureOptions{withMemory: tc.memory}, ) - machine := createExecutableBinding(t, ctx, fixture.Store, fixture.User.ID, name, fixture.Now.Add(time.Second)) + machine := createExecutableBinding(t, + ctx, + fixture.Store, + fixture.User.ID, + name, + fixture.Now.Add(time.Second)) var bindingID uuid.UUID err := fixture.Pool.QueryRow(ctx, ` INSERT INTO agent_machine_bindings( @@ -1330,14 +1350,22 @@ func TestFileToolsApprovalDispatch(t *testing.T) { if err != nil { t.Fatal(err) } - call := fixture.recordPendingToolCall(t, ctx, "call_file", name, string(input), fixture.Now.Add(20*time.Second)) + call := fixture.recordPendingToolCall(t, + ctx, + "call_file", + name, + string(input), + fixture.Now.Add(20*time.Second)) turn := fixture.turn() turn.Tools = map[string]ToolSpec{name: {Permission: toolpermission.DefaultSelection(toolpermission.ModeAlwaysAsk)}} wakes := 0 executor := Executor{ Store: fixture.Store, - - BackgroundRunner: backgroundRunnerFunc(func(string, func(context.Context) error) bool { wakes++; return true }), + BackgroundRunner: backgroundRunnerFunc(func(string, + func(context.Context) error) bool { + wakes++ + return true + }), } if err := executor.PrepareToolCallPermission(ctx, turn, call); err != nil { t.Fatal(err) @@ -1410,7 +1438,8 @@ func TestFileToolsApprovalDispatch(t *testing.T) { t.Fatalf("expected digest was not saved: %+v", target) } if _, err := fixture.Pool.Exec(ctx, - `UPDATE memory_stores SET deleted_at = statement_timestamp() WHERE id = $1`, target.StoreID); err != nil { + `UPDATE memory_stores SET deleted_at = statement_timestamp() WHERE id = $1`, + target.StoreID); err != nil { t.Fatal(err) } if _, err := fixture.Pool.Exec(ctx, @@ -1419,9 +1448,18 @@ SELECT project_id, name FROM memory_stores WHERE id = $1`, target.StoreID); err t.Fatal(err) } execution, err := fixture.Store.Execution().ExecuteToolCall(ctx, executionstore.ExecuteToolCallInput{ - ProjectID: toolsTestProjectID, AgentID: fixture.Agent.ID, ToolCallID: callID, RuntimeLockID: fixture.Lock.ID, + ProjectID: toolsTestProjectID, + AgentID: fixture.Agent.ID, + ToolCallID: callID, + RuntimeLockID: fixture.Lock.ID, }, func(reader *executionstore.ToolCallReader) (executionstore.ToolCallCommand, error) { - prepared, err := fileTransferProcessInput(ctx, reader, direction, "report.pdf", serverPath, target.ExpectedDigest) + prepared, + err := fileTransferProcessInput(ctx, + reader, + direction, + "report.pdf", + serverPath, + target.ExpectedDigest) if err != nil { return nil, err } @@ -1432,7 +1470,9 @@ SELECT project_id, name FROM memory_stores WHERE id = $1`, target.StoreID); err t.Fatal(err) } replay, ok := execution.CommandResult.(executionstore.ProcessRecord) - if !ok || replay.ID != process.ID || replay.ExecutionSpec.FileTransfer.Target.Memory.StoreID != target.StoreID { + if !ok || + replay.ID != process.ID || + replay.ExecutionSpec.FileTransfer.Target.Memory.StoreID != target.StoreID { t.Fatalf("replay changed the saved memory target: %+v", execution) } } @@ -1475,7 +1515,12 @@ func TestFileToolsUnavailableMemoryStoreFailsDispatch(t *testing.T) { "path": "/memory/" + storeName + "/report.pdf", tool.localField: "report.pdf", }) require.NoError(t, err) - call := fixture.recordToolCall(t, ctx, "call_file", tool.name, string(input), fixture.Now.Add(20*time.Second)) + call := fixture.recordToolCall(t, + ctx, + "call_file", + tool.name, + string(input), + fixture.Now.Add(20*time.Second)) turn := fixture.turn() turn.Tools[tool.name] = ToolSpec{Permission: toolpermission.DefaultSelection(toolpermission.ModeAlwaysAllow)} executor := Executor{Store: fixture.Store} @@ -1511,7 +1556,9 @@ func TestRemovedArtifactToolsFailDispatch(t *testing.T) { if err != nil { t.Fatal(err) } - if result.Disposition != DispatchCompleted || !strings.Contains(string(result.ContentParts), "unsupported") { + if result.Disposition != DispatchCompleted || + !strings.Contains(string(result.ContentParts), + "unsupported") { t.Fatalf("removed tool result=%+v", result) } }) diff --git a/internal/harness/tools/executor.go b/internal/harness/tools/executor.go index f356e5ed6..b5adaa655 100644 --- a/internal/harness/tools/executor.go +++ b/internal/harness/tools/executor.go @@ -20,6 +20,18 @@ func (e Executor) Dispatch(ctx context.Context, turn Turn, call model.ToolCall) if err != nil { return Result{}, fmt.Errorf("load tool call %s: %w", call.ID, err) } + return e.DispatchRecordedToolCall(ctx, turn, proposal) +} + +func (e Executor) DispatchRecordedToolCall( + ctx context.Context, turn Turn, proposal executionstore.ToolCallRecord, +) (Result, error) { + if proposal.ProjectID != turn.ProjectID || proposal.AgentID != turn.AgentID || + proposal.ModelCallContextID != turn.ModelCallContextID { + return Result{}, storeerr.ErrIdempotencyConflict + } + call := model.ToolCall{ID: proposal.ProviderCallID, Name: proposal.Name, Input: proposal.Input} + var err error if proposal.State == executionstore.ToolCallStateCompleted { return Result{ ToolCallID: call.ID, @@ -80,15 +92,10 @@ func (e Executor) Dispatch(ctx context.Context, turn Turn, call model.ToolCall) Disposition: DispatchDeferred, }, nil case toolDispatchCompleted: - if replayed, ok, err := e.completedToolResult(ctx, turn, call); err != nil { - return Result{}, err - } else if ok { - return replayed, nil - } - return Result{}, fmt.Errorf( - "tool %s completed during dispatch but no result is available", - call.Name, - ) + return Result{ + ToolCallID: call.ID, Name: call.Name, ContentParts: dispatchResult.record.ResultContentParts, + Disposition: DispatchCompleted, + }, nil case toolDispatchFailed: content = dispatchResult.content execErr = dispatchResult.cause @@ -250,30 +257,6 @@ func machineUnavailableToolResult(cause error) (machineUnavailableResult, error) return machineUnavailableResult{Cause: errors.New(errorCode), Content: content}, nil } -func (e Executor) completedToolResult( - ctx context.Context, - turn Turn, - call model.ToolCall, -) (Result, bool, error) { - record, err := e.recordedToolCall(ctx, turn, call) - if err != nil { - return Result{}, false, err - } - if !toolCallMatches(record, call) { - return Result{}, false, storeerr.ErrIdempotencyConflict - } - if record.State != executionstore.ToolCallStateCompleted { - return Result{}, false, nil - } - parts := record.ResultContentParts - return Result{ - ToolCallID: call.ID, - Name: call.Name, - ContentParts: parts, - Disposition: DispatchCompleted, - }, true, nil -} - func (e Executor) recordedToolCall( ctx context.Context, turn Turn, diff --git a/internal/harness/tools/model_context_test_helpers_test.go b/internal/harness/tools/model_context_test_helpers_test.go index 94e3fde70..8ea29f504 100644 --- a/internal/harness/tools/model_context_test_helpers_test.go +++ b/internal/harness/tools/model_context_test_helpers_test.go @@ -39,16 +39,17 @@ func claimNormalModelCallForToolsTest( } sourceModelOutputID = output.ID } - claim, err := store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtime.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agentConfigID, - InputEventSequence: inputEventSequence, - SourceModelCallContextID: sourceModelCallContextID, - SourceModelOutputID: sourceModelOutputID, - }) + prepared1, err := store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: projectID, + AgentID: agentID, + RuntimeLockID: runtime.ID, + OpeningInputIDs: openingInputIDs, + SourceModelCallContextID: sourceModelCallContextID, + SourceModelOutputID: sourceModelOutputID, + }) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim normal model call: %v", err) } diff --git a/internal/harness/worker/continuation.go b/internal/harness/worker/continuation.go new file mode 100644 index 000000000..22309eddb --- /dev/null +++ b/internal/harness/worker/continuation.go @@ -0,0 +1,54 @@ +package worker + +import ( + "context" + "errors" + "time" + + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ContinuationOptions struct { + // MaxModelStarts counts additional model work steps after the initial claim. + // Zero disables continuation when these options are supplied explicitly. + MaxModelStarts int + MaxDuration time.Duration +} + +func continuationOptions(options *ContinuationOptions) ContinuationOptions { + if options == nil { + return ContinuationOptions{MaxModelStarts: 2, MaxDuration: 30 * time.Second} + } + result := *options + if result.MaxDuration <= 0 { + result.MaxDuration = 30 * time.Second + } + return result +} + +func (options ContinuationOptions) canAdvance(elapsed time.Duration) bool { + return options.MaxModelStarts > 0 && elapsed < options.MaxDuration +} + +func runtimeExecutionError(err error, active *activeRuntime) error { + if errors.Is(err, context.Canceled) && active.cancelRequested.Load() { + return nil + } + if errors.Is(err, context.Canceled) && active.renewalFailed.Load() { + return storeerr.ErrRuntimeLockInactive + } + return err +} + +func (w *Worker) advanceOwnedWork( + ctx context.Context, + input executionstore.AdvanceOwnedAgentWorkInput, +) (claim executionstore.ClaimedAgentWork, continued bool, err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = panicError("advance owned agent work", recovered) + } + }() + return w.store.AdvanceOwnedAgentWork(ctx, input) +} diff --git a/internal/harness/worker/continuation_integration_test.go b/internal/harness/worker/continuation_integration_test.go new file mode 100644 index 000000000..5113e75be --- /dev/null +++ b/internal/harness/worker/continuation_integration_test.go @@ -0,0 +1,210 @@ +//go:build integration + +package worker + +import ( + "context" + "encoding/json" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/harness/kernel" + "github.com/omnara-ai/omnara/internal/harness/tools" + "github.com/omnara-ai/omnara/internal/model" + "github.com/omnara-ai/omnara/internal/modelcontext" + "github.com/omnara-ai/omnara/internal/storage" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/testutil/modeltest" + "github.com/stretchr/testify/require" +) + +type observedContinuationExecutor struct { + kernel.AgentExecutor + models []kernel.ModelWorkExecution + tools []kernel.ToolWorkExecution + afterModel func(context.Context) error + asyncGate chan struct{} +} + +func (e *observedContinuationExecutor) ExecuteModelWork(ctx context.Context, input kernel.ModelWorkExecution) error { + e.models = append(e.models, input) + if err := e.AgentExecutor.ExecuteModelWork(ctx, input); err != nil { + return err + } + if e.afterModel != nil { + return e.afterModel(ctx) + } + return nil +} + +func (e *observedContinuationExecutor) ExecuteModelWorkAndAdvance( + ctx context.Context, input kernel.ModelWorkExecution, _ kernel.ModelWorkAdvanceOptions, +) (executionstore.OwnedAgentWorkTransition, bool, error) { + return executionstore.OwnedAgentWorkTransition{}, false, e.ExecuteModelWork(ctx, input) +} + +func (e *observedContinuationExecutor) ExecuteToolWork(ctx context.Context, input kernel.ToolWorkExecution) error { + e.tools = append(e.tools, input) + if e.asyncGate == nil { + return e.AgentExecutor.ExecuteToolWork(ctx, input) + } + reservation, err := tools.ReserveAsyncExecution(ctx) + if err != nil { + return err + } + reservation.Start() + go func() { + select { + case <-e.asyncGate: + reservation.Done(e.AgentExecutor.ExecuteToolWork(ctx, input)) + case <-ctx.Done(): + reservation.Done(ctx.Err()) + } + }() + return nil +} + +func continuationToolResponse(index int) model.Response { + return model.Response{ + ID: fmt.Sprintf("response_%d", index), StopReason: model.StopReasonToolUse, + Content: modeltest.ResponsePartsForToolCalls([]model.ToolCall{{ + ID: fmt.Sprintf("call_%d", index), Name: "list_agents", Input: json.RawMessage(`{}`), + }}), + } +} + +func continuationFinalResponse(index int) model.Response { + return model.Response{ + ID: fmt.Sprintf("response_%d", index), StopReason: model.StopReasonEndTurn, + Content: []model.ResponsePart{{Type: "text", Text: "done"}}, + } +} + +func newObservedContinuationExecutor(store *storage.Store, responses ...model.Response) *observedContinuationExecutor { + client := &sequenceModelClient{providerModelSlug: "worker-kernel-test", responses: responses} + return &observedContinuationExecutor{AgentExecutor: kernel.AgentExecutor{ + Store: store, + ContextBuilder: modelcontext.Builder{ + Store: modelcontext.NewStore(store.Execution(), store.Artifacts(), store.Integrations()), + }, + ModelResolver: liveWorkerTestModelResolver(store, client), + }} +} + +func TestWorkerOwnedContinuationRunsKernelAndYieldsAtModelBudget(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agentID, userID := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + inputID := createWorkerInput(t, ctx, store, agentID, userID, "keep going", time.Now()) + executor := newObservedContinuationExecutor(store, + continuationToolResponse(1), continuationToolResponse(2), continuationToolResponse(3), continuationFinalResponse(4)) + worker := NewWorker(store.Execution(), executor, Options{}) + requireWorkerClaim(t, ctx, worker, "run owned chain") + require.Len(t, executor.models, 3) + require.Len(t, executor.tools, 3) + runtimeID := executor.models[0].RuntimeLockID + for _, step := range executor.models { + require.Equal(t, runtimeID, step.RuntimeLockID) + require.Equal(t, []uuid.UUID{inputID}, step.InputIDs) + } + for _, step := range executor.tools { + require.Equal(t, runtimeID, step.RuntimeLockID) + } + waitNoRuntimeLock(t, ctx, pool, agentID) + var wakeups int + require.NoError(t, pool.QueryRow(ctx, + `SELECT count(*) FROM agent_wakeups WHERE agent_id=$1 AND ready_at <= now()`, agentID, + ).Scan(&wakeups)) + require.Equal(t, 1, wakeups) + requireWorkerClaim(t, ctx, worker, "claim yielded continuation") + require.Len(t, executor.models, 4) + require.NotEqual(t, runtimeID, executor.models[3].RuntimeLockID) + waitNoRuntimeLock(t, ctx, pool, agentID) + waitNoWorkerWakeup(t, ctx, pool, agentID) + assertCompletedToolCallCount(t, ctx, pool, agentID, "list_agents", "completed", 3) +} + +func TestWorkerOwnedContinuationAdmitsInputsBetweenSteps(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agentID, userID := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + initial := createWorkerInput(t, ctx, store, agentID, userID, "initial", time.Now()) + executor := newObservedContinuationExecutor(store, + continuationToolResponse(1), continuationFinalResponse(2), continuationFinalResponse(3)) + var steering, queued uuid.UUID + executor.afterModel = func(ctx context.Context) error { + if len(executor.models) == 1 { + queued = createWorkerInput(t, ctx, store, agentID, userID, "queued", time.Now()) + steering = createWorkerInputWithDeliveryMode( + t, ctx, store, agentID, userID, "steering", executionstore.DeliveryModeSteering, time.Now(), + ) + } + return nil + } + worker := NewWorker(store.Execution(), executor, Options{}) + requireWorkerClaim(t, ctx, worker, "admit arriving inputs") + require.Len(t, executor.models, 3) + require.Len(t, executor.tools, 1) + for i, inputID := range []uuid.UUID{initial, steering, queued} { + require.Equal(t, []uuid.UUID{inputID}, executor.models[i].InputIDs) + require.Equal(t, executor.models[0].RuntimeLockID, executor.models[i].RuntimeLockID) + } + require.Equal(t, executor.models[0].TurnID, executor.tools[0].TurnID) + waitNoRuntimeLock(t, ctx, pool, agentID) + waitNoWorkerWakeup(t, ctx, pool, agentID) +} + +func TestWorkerOwnedContinuationHonorsDurableCancellationBetweenSteps(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agentID, userID := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + createWorkerInput(t, ctx, store, agentID, userID, "cancel", time.Now()) + executor := newObservedContinuationExecutor(store, continuationToolResponse(1)) + executor.afterModel = func(ctx context.Context) error { + _, err := store.Execution().CancelAgent(ctx, executionstore.CancelAgentInput{ + ProjectID: workerTestProjectID, AgentID: agentID, Actor: workerTestActorParams(t), + }) + return err + } + worker := NewWorker(store.Execution(), executor, Options{}) + worked, err := worker.RunOnce(ctx) + require.True(t, worked) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) + require.Len(t, executor.models, 1) + require.Empty(t, executor.tools) + waitNoRuntimeLock(t, ctx, pool, agentID) + waitNoWorkerWakeup(t, ctx, pool, agentID) +} + +func TestWorkerOwnedContinuationTransfersLeaseToAsyncRetention(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agentID, userID := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + createWorkerInput(t, ctx, store, agentID, userID, "async", time.Now()) + executor := newObservedContinuationExecutor(store, continuationToolResponse(1), continuationFinalResponse(2)) + executor.asyncGate = make(chan struct{}) + worker := NewWorker(store.Execution(), executor, Options{AsyncToolCapacity: 1}) + requireWorkerClaim(t, ctx, worker, "retain async continuation") + require.Len(t, executor.models, 1) + require.Len(t, executor.tools, 1) + var runtimeID uuid.UUID + require.NoError(t, pool.QueryRow(ctx, + `SELECT id FROM agent_runtime_locks WHERE agent_id=$1`, agentID, + ).Scan(&runtimeID)) + require.Equal(t, executor.models[0].RuntimeLockID, runtimeID) + require.Equal(t, runtimeID, executor.tools[0].RuntimeLockID) + close(executor.asyncGate) + worker.retained.Wait() + waitNoRuntimeLock(t, ctx, pool, agentID) + requireWorkerClaim(t, ctx, worker, "resume after async completion") + require.Len(t, executor.models, 2) + require.NotEqual(t, runtimeID, executor.models[1].RuntimeLockID) + waitNoWorkerWakeup(t, ctx, pool, agentID) +} diff --git a/internal/harness/worker/continuation_test.go b/internal/harness/worker/continuation_test.go new file mode 100644 index 000000000..d2fbc90db --- /dev/null +++ b/internal/harness/worker/continuation_test.go @@ -0,0 +1,266 @@ +package worker + +import ( + "context" + "errors" + "reflect" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/harness/kernel" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type continuationStore struct { + retainedRuntimeStore + next []executionstore.ClaimedAgentWork + advances []executionstore.AdvanceOwnedAgentWorkInput + advanceError error + advancePanic bool + advanceDelay time.Duration +} + +func (s *continuationStore) AdvanceOwnedAgentWork( + ctx context.Context, + input executionstore.AdvanceOwnedAgentWorkInput, +) (executionstore.ClaimedAgentWork, bool, error) { + s.advances = append(s.advances, input) + if s.advanceDelay > 0 { + time.Sleep(s.advanceDelay) //nolint:omnaralint // Exercise elapsed time; no concurrent state is being awaited. + } + if s.advancePanic { + panic("advance failed") + } + if s.advanceError != nil { + return executionstore.ClaimedAgentWork{}, false, s.advanceError + } + if len(s.next) == 0 || (s.next[0].Kind == executionstore.AgentWorkModel && !input.AllowModelWork) { + err := s.ReleaseAgentRuntimeLock(ctx, input.ProjectID, input.AgentID, input.RuntimeLockID) + return executionstore.ClaimedAgentWork{}, false, err + } + claim := s.next[0] + s.next = s.next[1:] + return claim, true, nil +} + +type continuationExecutor struct { + kinds []executionstore.AgentWorkKind + runtimes []uuid.UUID + afterStep func(context.Context) error +} + +func (e *continuationExecutor) execute( + ctx context.Context, + kind executionstore.AgentWorkKind, + runtime uuid.UUID, +) error { + e.kinds = append(e.kinds, kind) + e.runtimes = append(e.runtimes, runtime) + if e.afterStep != nil { + return e.afterStep(ctx) + } + return nil +} + +func (e *continuationExecutor) ExecuteModelWork( + ctx context.Context, + input kernel.ModelWorkExecution, +) error { + return e.execute(ctx, executionstore.AgentWorkModel, input.RuntimeLockID) +} + +func (e *continuationExecutor) ExecuteToolWork( + ctx context.Context, + input kernel.ToolWorkExecution, +) error { + return e.execute(ctx, executionstore.AgentWorkTool, input.RuntimeLockID) +} + +func newContinuationStore(kinds ...executionstore.AgentWorkKind) *continuationStore { + claim := executionstore.ClaimedAgentWork{ + OrgID: uuid.New(), ProjectID: uuid.New(), AgentID: uuid.New(), + Kind: executionstore.AgentWorkModel, + RuntimeLock: executionstore.AgentRuntimeLockRecord{ID: uuid.New()}, + } + s := &continuationStore{retainedRuntimeStore: retainedRuntimeStore{claim: claim, released: make(chan uuid.UUID, 8)}} + for _, kind := range kinds { + next := claim + next.Kind = kind + s.next = append(s.next, next) + } + return s +} + +func TestWorkerOwnedContinuationBudget(t *testing.T) { + model, tool := executionstore.AgentWorkModel, executionstore.AgentWorkTool + for _, test := range []struct { + name string + options *ContinuationOptions + want []executionstore.AgentWorkKind + }{ + {"default", nil, []executionstore.AgentWorkKind{model, tool, model, tool, model, tool}}, + { + "one additional model", &ContinuationOptions{MaxModelStarts: 1}, + []executionstore.AgentWorkKind{model, tool, model, tool}, + }, + {"disabled", &ContinuationOptions{}, []executionstore.AgentWorkKind{model}}, + { + "elapsed", &ContinuationOptions{MaxModelStarts: 5, MaxDuration: time.Nanosecond}, + []executionstore.AgentWorkKind{model}, + }, + } { + t.Run(test.name, func(t *testing.T) { + store := newContinuationStore(tool, model, tool, model, tool, model) + executor := &continuationExecutor{} + worker := NewWorker(store, executor, Options{Continuation: test.options}) + worked, err := worker.RunOnce(t.Context()) + if err != nil || !worked { + t.Fatalf("run: worked=%v err=%v", worked, err) + } + if !reflect.DeepEqual(executor.kinds, test.want) { + t.Fatalf("steps=%v want=%v", executor.kinds, test.want) + } + for _, runtime := range executor.runtimes { + if runtime != store.claim.RuntimeLock.ID { + t.Fatalf("runtime changed: %s", runtime) + } + } + if store.renewalCalls.Load() != 1 { + t.Fatalf("initial renewals=%d want=1", store.renewalCalls.Load()) + } + if len(store.released) != 1 { + t.Fatalf("releases=%d want=1", len(store.released)) + } + if runtime := <-store.released; runtime != store.claim.RuntimeLock.ID { + t.Fatalf("released %s", runtime) + } + }) + } +} + +func TestWorkerOwnedContinuationFailureReleasesOriginalRuntime(t *testing.T) { + for _, test := range []struct { + name string + err error + panic bool + }{ + {"store error", errors.New("database unavailable"), false}, + {"inactive lease", storeerr.ErrRuntimeLockInactive, false}, + {"panic", nil, true}, + } { + t.Run(test.name, func(t *testing.T) { + store := newContinuationStore() + store.advanceError, store.advancePanic = test.err, test.panic + worker := NewWorker(store, &continuationExecutor{}, Options{}) + worked, err := worker.RunOnce(t.Context()) + if !worked || err == nil { + t.Fatalf("worked=%v err=%v", worked, err) + } + if test.err != nil && !errors.Is(err, test.err) { + t.Fatalf("error=%v want=%v", err, test.err) + } + if len(store.released) != 1 || <-store.released != store.claim.RuntimeLock.ID { + t.Fatal("original runtime was not released once") + } + if len(worker.active) != 0 { + t.Fatal("runtime still registered") + } + }) + } +} + +func TestWorkerOwnedContinuationCancellationBetweenSteps(t *testing.T) { + store := newContinuationStore(executionstore.AgentWorkTool) + executor := &continuationExecutor{} + worker := NewWorker(store, executor, Options{}) + executor.afterStep = func(context.Context) error { + worker.activeMu.Lock() + active := worker.active[store.claim.RuntimeLock.ID] + worker.activeMu.Unlock() + active.cancelRequested.Store(true) + active.cancel() + return nil + } + worked, err := worker.RunOnce(t.Context()) + if err != nil || !worked { + t.Fatalf("worked=%v err=%v", worked, err) + } + if len(store.advances) != 0 || len(store.released) != 1 { + t.Fatalf("advances=%d releases=%d", len(store.advances), len(store.released)) + } +} + +func TestWorkerOwnedContinuationKeepsRenewalRunning(t *testing.T) { + store := newContinuationStore(executionstore.AgentWorkModel) + store.renewed = make(chan struct{}) + executor := &continuationExecutor{} + executor.afterStep = func(ctx context.Context) error { + if len(executor.kinds) == 2 { + select { + case <-store.renewed: + case <-ctx.Done(): + return ctx.Err() + case <-time.After(5 * time.Second): + return errors.New("renewal stopped between steps") + } + } + return nil + } + worker := NewWorker(store, executor, Options{RuntimeLockLeaseDuration: 150 * time.Millisecond}) + worked, err := worker.RunOnce(t.Context()) + if err != nil || !worked { + t.Fatalf("worked=%v err=%v", worked, err) + } + if len(executor.kinds) != 2 || len(store.released) != 1 { + t.Fatalf("steps=%v releases=%d", executor.kinds, len(store.released)) + } +} + +func TestWorkerOwnedContinuationCountsTimeSpentAdvancing(t *testing.T) { + store := newContinuationStore(executionstore.AgentWorkModel) + store.advanceDelay = 110 * time.Millisecond + executor := &continuationExecutor{} + worker := NewWorker(store, executor, Options{ + Continuation: &ContinuationOptions{MaxModelStarts: 2, MaxDuration: 100 * time.Millisecond}, + }) + worked, err := worker.RunOnce(t.Context()) + if err != nil || !worked { + t.Fatalf("worked=%v err=%v", worked, err) + } + if len(executor.kinds) != 1 || len(store.advances) != 1 || len(store.released) != 1 { + t.Fatalf("steps=%v advances=%d releases=%d", executor.kinds, len(store.advances), len(store.released)) + } +} + +func TestWorkerUsesRemainingClaimLeaseBudget(t *testing.T) { + for _, test := range []struct { + name string + age time.Duration + wantRenewals int32 + }{ + {"fresh", 0, 0}, {"delayed beyond local cutoff", time.Minute, 1}, + } { + t.Run(test.name, func(t *testing.T) { + store := newContinuationStore() + store.claim.LocalLeaseBudgetStartedAt = time.Now().Add(-test.age) + worker := NewWorker(store, &fusedContinuationTestExecutor{}, Options{RuntimeLockLeaseDuration: time.Minute}) + worked, err := worker.RunOnce(t.Context()) + if err != nil || !worked { + t.Fatalf("worked=%v err=%v", worked, err) + } + if got := store.renewalCalls.Load(); got != test.wantRenewals { + t.Fatalf("renewals=%d want=%d", got, test.wantRenewals) + } + }) + } +} + +type fusedContinuationTestExecutor struct{ continuationExecutor } + +func (e *fusedContinuationTestExecutor) ExecuteModelWorkAndAdvance( + ctx context.Context, input kernel.ModelWorkExecution, _ kernel.ModelWorkAdvanceOptions, +) (executionstore.OwnedAgentWorkTransition, bool, error) { + return executionstore.OwnedAgentWorkTransition{}, false, e.ExecuteModelWork(ctx, input) +} diff --git a/internal/harness/worker/round_trips_integration_test.go b/internal/harness/worker/round_trips_integration_test.go new file mode 100644 index 000000000..151555c61 --- /dev/null +++ b/internal/harness/worker/round_trips_integration_test.go @@ -0,0 +1,106 @@ +//go:build integration + +package worker + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/multitracer" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/omnara-ai/omnara/internal/model" + "github.com/omnara-ai/omnara/internal/modelprovider" + "github.com/omnara-ai/omnara/internal/storage" + "github.com/stretchr/testify/require" +) + +func TestWorkerModelToolChainPoolAcquisitions(t *testing.T) { + ctx := t.Context() + fixture := openWorkerIntegrationDB(t, ctx) + cfg := fixture.Config() + trace := &stepCommitTracer{} + cfg.ConnConfig.Tracer = multitracer.New(cfg.ConnConfig.Tracer, trace) + pool, err := pgxpool.NewWithConfig(ctx, cfg) + require.NoError(t, err) + defer pool.Close() + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agentID, userID := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + createWorkerInput(t, ctx, store, agentID, userID, "continue", time.Now()) + executor := newObservedContinuationExecutor(store, + continuationToolResponse(1), continuationToolResponse(2), continuationFinalResponse(3)) + executor.ModelResolver = credentialResolvingTestModel{ + resolver: modelprovider.Resolver{Models: store.Models(), Secrets: store.Secrets()}, + client: &sequenceModelClient{providerModelSlug: "worker-kernel-test", responses: []model.Response{ + continuationToolResponse(1), continuationToolResponse(2), continuationFinalResponse(3), + }}, + } + worker := NewWorker(store.Execution(), executor.AgentExecutor, Options{}) + before := pool.Stat().AcquireCount() + commitsBefore := trace.commits.Load() + requireWorkerClaim(t, ctx, worker, "model/tool chain") + acquires := pool.Stat().AcquireCount() - before + commits := trace.commits.Load() - commitsBefore + t.Logf("three model calls and two tools: %d pool acquisitions, %d commits", acquires, commits) + waitNoRuntimeLock(t, ctx, pool, agentID) + waitNoWorkerWakeup(t, ctx, pool, agentID) + assertCompletedToolCallCount(t, ctx, pool, agentID, "list_agents", "completed", 2) + require.LessOrEqual(t, acquires, int64(30)) + require.LessOrEqual(t, commits, int64(12)) +} + +type stepCommitTracer struct{ commits atomic.Int64 } + +func (*stepCommitTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData) context.Context { + return ctx +} +func (t *stepCommitTracer) TraceQueryEnd(_ context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) { + if data.Err == nil && data.CommandTag.String() == "COMMIT" { + t.commits.Add(1) + } +} + +type credentialResolvingTestModel struct { + resolver modelprovider.Resolver + client model.Client +} + +func (r credentialResolvingTestModel) Resolve( + ctx context.Context, selection model.Selection, +) (model.ResolvedClient, error) { + resolved, err := r.resolver.Resolve(ctx, selection) + if err != nil { + return model.ResolvedClient{}, err + } + resolved.Client = r.client + return resolved, nil +} + +func TestWorkerFusedContinuationYieldsBeforePreparingOverBudgetModel(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agentID, userID := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + createWorkerInput(t, ctx, store, agentID, userID, "continue", time.Now()) + executor := newObservedContinuationExecutor(store) + client := &sequenceModelClient{providerModelSlug: "worker-kernel-test", responses: []model.Response{ + continuationToolResponse(1), continuationToolResponse(2), continuationToolResponse(3), continuationFinalResponse(4), + }} + executor.ModelResolver = liveWorkerTestModelResolver(store, client) + worker := NewWorker(store.Execution(), executor.AgentExecutor, Options{}) + requireWorkerClaim(t, ctx, worker, "bounded fused chain") + require.Equal(t, 3, client.calls) + assertCompletedToolCallCount(t, ctx, pool, agentID, "list_agents", "completed", 3) + waitNoRuntimeLock(t, ctx, pool, agentID) + var contexts int + require.NoError(t, pool.QueryRow(ctx, + `SELECT count(*) FROM model_call_contexts WHERE agent_id=$1`, agentID, + ).Scan(&contexts)) + require.Equal(t, 3, contexts, "yielding must not prepare an unused fourth attempt") + requireWorkerClaim(t, ctx, worker, "reclaim yielded chain") + require.Equal(t, 4, client.calls) + waitNoWorkerWakeup(t, ctx, pool, agentID) + waitNoRuntimeLock(t, ctx, pool, agentID) +} diff --git a/internal/harness/worker/worker.go b/internal/harness/worker/worker.go index 44ae78ae9..d21eecd72 100644 --- a/internal/harness/worker/worker.go +++ b/internal/harness/worker/worker.go @@ -27,11 +27,21 @@ type AgentWorkExecutor interface { ExecuteToolWork(context.Context, kernel.ToolWorkExecution) error } +type modelWorkAdvancer interface { + ExecuteModelWorkAndAdvance( + context.Context, kernel.ModelWorkExecution, kernel.ModelWorkAdvanceOptions, + ) (executionstore.OwnedAgentWorkTransition, bool, error) +} + type Store interface { ClaimNextAgentWork( context.Context, executionstore.ClaimNextAgentWorkInput, ) (executionstore.ClaimedAgentWork, bool, error) + AdvanceOwnedAgentWork( + context.Context, + executionstore.AdvanceOwnedAgentWorkInput, + ) (executionstore.ClaimedAgentWork, bool, error) ReleaseAgentRuntimeLock(context.Context, uuid.UUID, uuid.UUID, uuid.UUID) error RenewAgentRuntimeLock( context.Context, @@ -50,18 +60,23 @@ type Worker struct { controlSubscriber notifications.WorkerControlSubscriber runtimeLockLeaseDuration time.Duration capacity int + claimConcurrency int + executionFailures atomic.Int64 asyncToolLimiter *tools.AsyncExecutionLimiter activeMu sync.Mutex active map[uuid.UUID]*activeRuntime retained sync.WaitGroup + continuation ContinuationOptions } type Options struct { Log *slog.Logger RuntimeLockLeaseDuration time.Duration Capacity int + ClaimConcurrency int AsyncToolCapacity int ControlSubscriber notifications.WorkerControlSubscriber + Continuation *ContinuationOptions } type activeRuntime struct { @@ -73,6 +88,9 @@ type activeRuntime struct { const ( defaultAsyncToolCapacity = 32 + defaultClaimConcurrency = 4 + minIdleClaimDelay = 50 * time.Millisecond + maxIdleClaimDelay = time.Second initialRunLoopRetryDelay = 250 * time.Millisecond maxRunLoopRetryDelay = 30 * time.Second ) @@ -90,10 +108,14 @@ func NewWorker(store Store, executor AgentWorkExecutor, opts Options) *Worker { if opts.Capacity <= 0 { opts.Capacity = 4 } + if opts.ClaimConcurrency <= 0 { + opts.ClaimConcurrency = defaultClaimConcurrency + } if opts.AsyncToolCapacity <= 0 { opts.AsyncToolCapacity = defaultAsyncToolCapacity } return &Worker{ + continuation: continuationOptions(opts.Continuation), store: store, executor: executor, log: log, @@ -101,6 +123,7 @@ func NewWorker(store Store, executor AgentWorkExecutor, opts Options) *Worker { controlSubscriber: opts.ControlSubscriber, runtimeLockLeaseDuration: opts.RuntimeLockLeaseDuration, capacity: opts.Capacity, + claimConcurrency: min(opts.ClaimConcurrency, opts.Capacity), asyncToolLimiter: tools.NewAsyncExecutionLimiter(opts.AsyncToolCapacity), active: make(map[uuid.UUID]*activeRuntime), } @@ -113,16 +136,18 @@ func (w *Worker) Run(ctx context.Context) error { return err } defer func() { _ = controlSubscription.Unsubscribe() }() - var loops sync.WaitGroup - for range w.capacity { - loops.Add(1) + slots := make(chan struct{}, w.capacity) + var claimers, executions sync.WaitGroup + for range w.claimConcurrency { + claimers.Add(1) go func() { - defer loops.Done() - w.runLoop(ctx) + defer claimers.Done() + w.claimLoop(ctx, slots, &executions) }() } <-ctx.Done() - loops.Wait() + claimers.Wait() + executions.Wait() w.retained.Wait() return ctx.Err() } @@ -134,67 +159,103 @@ func (w *Worker) subscribeWorkerControl(ctx context.Context) (notifications.Subs return w.controlSubscriber.SubscribeWorkerControl(ctx, w.workerProcessID, w.handleWorkerControl) } -func (w *Worker) runLoop(ctx context.Context) { - idleDelay := 50 * time.Millisecond +func (w *Worker) claimLoop(ctx context.Context, slots chan struct{}, executions *sync.WaitGroup) { + idleDelay := minIdleClaimDelay failures := 0 for { - loopCtx, event := logent.WorkerLoop(ctx, w.workerProcessID) - worked, err := w.RunOnce(loopCtx) - logent.WorkerLoopResult(loopCtx, worked, err) - recoverable := isRecoverableRunOnceError(err) - if recoverable { - logent.WorkerLoopRecoverableTurnRace(loopCtx, err) + select { + case slots <- struct{}{}: + case <-ctx.Done(): + return } - switch { - case err != nil && !recoverable && - !(ctx.Err() != nil && errutil.OnlyMatches(err, context.Canceled)): - logpkg.Error(loopCtx, err) + loopCtx, event := logent.WorkerLoop(ctx, w.workerProcessID) + now := time.Now().UTC() + claim, handled, err := w.claimNext(loopCtx) + if err == nil && handled && claim.Kind != executionstore.AgentWorkNone { + failures = 0 + idleDelay = minIdleClaimDelay + executions.Add(1) + go func() { + defer executions.Done() + defer func() { <-slots }() + err := w.executeClaimedWorkRecovering(loopCtx, claim, now) + w.finishLoop(ctx, loopCtx, event, true, err) + w.throttleAfterExecution(ctx, err) + }() + continue } - event.Done(loopCtx) + <-slots + w.finishLoop(ctx, loopCtx, event, handled, err) if ctx.Err() != nil { return } - if recoverable { + var delay time.Duration + switch { + case isRecoverableRunOnceError(err): failures = 0 - idleDelay = 50 * time.Millisecond - if !waitForDelay(ctx, initialRunLoopRetryDelay) { - return - } - continue - } - if err != nil { - delay := runLoopRetryDelay(failures) + idleDelay = minIdleClaimDelay + delay = initialRunLoopRetryDelay + case err != nil: + delay = runLoopRetryDelay(failures) failures++ - if !waitForDelay(ctx, delay) { - return - } - continue - } - failures = 0 - if worked { - idleDelay = 50 * time.Millisecond + case handled: + failures = 0 + idleDelay = minIdleClaimDelay continue + default: + failures = 0 + delay = idleDelay + idleDelay = min(idleDelay*2, maxIdleClaimDelay) } - if !waitForDelay(ctx, idleDelay) { + if !waitForDelay(ctx, delay) { return } - if idleDelay < time.Second { - idleDelay *= 2 - if idleDelay > time.Second { - idleDelay = time.Second - } - } + } +} + +func (w *Worker) finishLoop(ctx, loopCtx context.Context, event *logpkg.Event, worked bool, err error) { + logent.WorkerLoopResult(loopCtx, worked, err) + recoverable := isRecoverableRunOnceError(err) + if recoverable { + logent.WorkerLoopRecoverableTurnRace(loopCtx, err) + } + if err != nil && !recoverable && !(ctx.Err() != nil && errutil.OnlyMatches(err, context.Canceled)) { + logpkg.Error(loopCtx, err) + } + event.Done(loopCtx) +} + +func (w *Worker) throttleAfterExecution(ctx context.Context, err error) { + switch { + case err == nil || ctx.Err() != nil: + w.executionFailures.Store(0) + case isRecoverableRunOnceError(err): + w.executionFailures.Store(0) + waitForDelay(ctx, initialRunLoopRetryDelay) + default: + waitForDelay(ctx, runLoopRetryDelay(int(w.executionFailures.Add(1)-1))) } } func (w *Worker) RunOnce(ctx context.Context) (worked bool, err error) { + now := time.Now().UTC() + claim, handled, err := w.claimNext(ctx) + if err != nil || !handled { + return false, err + } + if claim.Kind == executionstore.AgentWorkNone { + return true, nil + } + return true, w.executeClaimedWorkRecovering(ctx, claim, now) +} + +func (w *Worker) claimNext(ctx context.Context) (claim executionstore.ClaimedAgentWork, handled bool, err error) { defer func() { if recovered := recover(); recovered != nil { - err = panicError("worker iteration", recovered) + claim, handled, err = executionstore.ClaimedAgentWork{}, false, panicError("worker claim", recovered) } }() - now := time.Now().UTC() - claim, handled, err := w.store.ClaimNextAgentWork( + claim, handled, err = w.store.ClaimNextAgentWork( ctx, executionstore.ClaimNextAgentWorkInput{ WorkerProcessID: w.workerProcessID, @@ -202,18 +263,25 @@ func (w *Worker) RunOnce(ctx context.Context) (worked bool, err error) { }, ) if errors.Is(err, storeerr.ErrNoClaimableAgentWakeup) { - return false, nil + return executionstore.ClaimedAgentWork{}, false, nil } if err != nil { - return false, err - } - if !handled { - return false, nil + return executionstore.ClaimedAgentWork{}, false, err } - if claim.Kind == executionstore.AgentWorkNone { - return true, nil - } - return true, w.executeClaimedWork(ctx, claim, now) + return claim, handled, nil +} + +func (w *Worker) executeClaimedWorkRecovering( + ctx context.Context, + claim executionstore.ClaimedAgentWork, + now time.Time, +) (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = panicError("worker iteration", recovered) + } + }() + return w.executeClaimedWork(ctx, claim, now) } func (w *Worker) executeClaimedWork( @@ -222,25 +290,27 @@ func (w *Worker) executeClaimedWork( now time.Time, ) error { runtime := claim.RuntimeLock + projectID, agentID := claim.ProjectID, claim.AgentID logent.AgentWorkScope(ctx, claim.OrgID, claim.ProjectID, claim.AgentID) logent.RuntimeLock(ctx, runtime) - if claim.Kind == executionstore.AgentWorkModel && - claim.Model.AdmittedInputTurn.Turn.ID != uuid.Nil { - logent.AdmittedAgentInputTurn(ctx, claim.Model.AdmittedInputTurn) + leaseBudgetStartedAt := claim.LocalLeaseBudgetStartedAt + if _, fenced := w.executor.(modelWorkAdvancer); !fenced { + leaseBudgetStartedAt = time.Time{} } turnCtx, active, stopRenewal, err := w.startRuntimeRenewal( ctx, - claim.ProjectID, - claim.AgentID, + projectID, + agentID, runtime, + leaseBudgetStartedAt, ) if err != nil { finalizeCtx, cancelFinalize := w.finalizeContext(ctx) defer cancelFinalize() if releaseErr := w.store.ReleaseAgentRuntimeLock( finalizeCtx, - claim.ProjectID, - claim.AgentID, + projectID, + agentID, runtime.ID, ); releaseErr != nil && !errors.Is(releaseErr, storeerr.ErrRuntimeLockInactive) { return releaseErr @@ -250,38 +320,79 @@ func (w *Worker) executeClaimedWork( } return err } - asyncScope := tools.NewAsyncExecutionScope(w.asyncToolLimiter) - turnCtx = tools.WithAsyncExecutionScope(turnCtx, asyncScope) - err = w.executeWork(turnCtx, claim, now) - asyncScope.Seal() - if asyncScope.Started() { - if errors.Is(err, context.Canceled) && active.cancelRequested.Load() { - err = nil - } else if errors.Is(err, context.Canceled) && active.renewalFailed.Load() { - err = storeerr.ErrRuntimeLockInactive + startedAt := time.Now() + additionalModelStarts := 0 + for { + if claim.Kind == executionstore.AgentWorkModel && claim.Model.AdmittedInputTurn.Turn.ID != uuid.Nil { + logent.AdmittedAgentInputTurn(ctx, claim.Model.AdmittedInputTurn) } - w.retainRuntimeUntilAsyncCompletion( - ctx, - claim.ProjectID, - claim.AgentID, - runtime.ID, - asyncScope, - stopRenewal, + asyncScope := tools.NewAsyncExecutionScope(w.asyncToolLimiter) + var advance *kernel.ModelWorkAdvanceOptions + if w.continuation.canAdvance(time.Since(startedAt)) { + advance = &kernel.ModelWorkAdvanceOptions{ + Deadline: startedAt.Add(w.continuation.MaxDuration), + AllowModelWork: additionalModelStarts < w.continuation.MaxModelStarts, + } + } + transition, advanced, stepErr := w.executeWorkWithAdvance( + tools.WithAsyncExecutionScope(turnCtx, asyncScope), claim, now, advance, ) - return err + err = stepErr + asyncScope.Seal() + if asyncScope.Started() { + w.retainRuntimeUntilAsyncCompletion( + ctx, projectID, agentID, runtime.ID, asyncScope, stopRenewal, + ) + return runtimeExecutionError(err, active) + } + if err == nil && advanced && !transition.Continued { + stopRenewal() + return nil + } + if err == nil { + err = turnCtx.Err() + } + if err != nil { + break + } + if advanced { + claim = transition.Work + } else { + if !w.continuation.canAdvance(time.Since(startedAt)) { + break + } + var continued bool + _, prepare := w.executor.(modelWorkAdvancer) + claim, continued, err = w.advanceOwnedWork(turnCtx, executionstore.AdvanceOwnedAgentWorkInput{ + ProjectID: projectID, AgentID: agentID, RuntimeLockID: runtime.ID, + AllowModelWork: additionalModelStarts < w.continuation.MaxModelStarts, + PrepareModel: prepare, + }) + if err != nil { + break + } + if !continued { + stopRenewal() + return nil + } + } + err = turnCtx.Err() + if err != nil || !w.continuation.canAdvance(time.Since(startedAt)) { + break + } + if claim.Kind == executionstore.AgentWorkModel { + additionalModelStarts++ + } + now = time.Now().UTC() } stopRenewal() finalizeCtx, cancelFinalize := w.finalizeContext(ctx) defer cancelFinalize() - if errors.Is(err, context.Canceled) && active.cancelRequested.Load() { - err = nil - } else if errors.Is(err, context.Canceled) && active.renewalFailed.Load() { - err = storeerr.ErrRuntimeLockInactive - } + err = runtimeExecutionError(err, active) if releaseErr := w.store.ReleaseAgentRuntimeLock( finalizeCtx, - claim.ProjectID, - claim.AgentID, + projectID, + agentID, runtime.ID, ); releaseErr != nil && err == nil { @@ -327,6 +438,13 @@ func (w *Worker) executeWork( claim executionstore.ClaimedAgentWork, now time.Time, ) (err error) { + _, _, err = w.executeWorkWithAdvance(ctx, claim, now, nil) + return err +} + +func (w *Worker) executeWorkWithAdvance( + ctx context.Context, claim executionstore.ClaimedAgentWork, now time.Time, advance *kernel.ModelWorkAdvanceOptions, +) (transition executionstore.OwnedAgentWorkTransition, advanced bool, err error) { defer func() { if recovered := recover(); recovered != nil { err = panicError("agent work", recovered) @@ -334,7 +452,8 @@ func (w *Worker) executeWork( }() switch claim.Kind { case executionstore.AgentWorkModel: - return w.executor.ExecuteModelWork(ctx, kernel.ModelWorkExecution{ + input := kernel.ModelWorkExecution{ + Prepared: claim.Model.Prepared, Kind: claim.Model.Kind, OrgID: claim.OrgID, ProjectID: claim.ProjectID, @@ -347,9 +466,14 @@ func (w *Worker) executeWork( OpeningEventSequence: claim.Model.OpeningEventSequence, RuntimeLockID: claim.RuntimeLock.ID, Now: now, - }) + } + if executor, ok := w.executor.(modelWorkAdvancer); ok && advance != nil { + return executor.ExecuteModelWorkAndAdvance(ctx, input, *advance) + } + return transition, false, w.executor.ExecuteModelWork(ctx, input) case executionstore.AgentWorkTool: - return w.executor.ExecuteToolWork(ctx, kernel.ToolWorkExecution{ + return transition, false, w.executor.ExecuteToolWork(ctx, kernel.ToolWorkExecution{ + Prepared: claim.Tool.Prepared, ProjectID: claim.ProjectID, AgentID: claim.AgentID, TurnID: claim.Tool.TurnID, @@ -360,7 +484,7 @@ func (w *Worker) executeWork( Now: now, }) default: - return fmt.Errorf("unsupported agent work kind %d", claim.Kind) + return transition, false, fmt.Errorf("unsupported agent work kind %d", claim.Kind) } } @@ -419,6 +543,7 @@ func (w *Worker) startRuntimeRenewal( ctx context.Context, projectID, agentID uuid.UUID, runtime executionstore.AgentRuntimeLockRecord, + leaseBudgetStartedAt time.Time, ) (context.Context, *activeRuntime, func(), error) { turnCtx, cancelTurn := context.WithCancel(ctx) active, unregister := w.registerActiveRuntime( @@ -426,21 +551,25 @@ func (w *Worker) startRuntimeRenewal( runtime, cancelTurn, ) - renewalCutoff, err := w.renewRuntimeUntil( - turnCtx, - projectID, - agentID, - runtime.ID, - runtimeRenewalCutoff(time.Now(), w.runtimeLockLeaseDuration), - ) - if err != nil { - if errors.Is(err, errRuntimeCancelRequested) || - (errors.Is(err, context.Canceled) && active.cancelRequested.Load()) { - err = errRuntimeCancelRequested + renewalCutoff := runtimeRenewalCutoff(leaseBudgetStartedAt, w.runtimeLockLeaseDuration) + if leaseBudgetStartedAt.IsZero() || !time.Now().Before(renewalCutoff) { + var err error + renewalCutoff, err = w.renewRuntimeUntil( + turnCtx, + projectID, + agentID, + runtime.ID, + runtimeRenewalCutoff(time.Now(), w.runtimeLockLeaseDuration), + ) + if err != nil { + if errors.Is(err, errRuntimeCancelRequested) || + (errors.Is(err, context.Canceled) && active.cancelRequested.Load()) { + err = errRuntimeCancelRequested + } + unregister() + cancelTurn() + return nil, nil, nil, err } - unregister() - cancelTurn() - return nil, nil, nil, err } renewalCtx, cancelRenewal := context.WithCancel(turnCtx) done := make(chan struct{}) diff --git a/internal/harness/worker/worker_integration_test.go b/internal/harness/worker/worker_integration_test.go index dc1c8323e..5e2620eb3 100644 --- a/internal/harness/worker/worker_integration_test.go +++ b/internal/harness/worker/worker_integration_test.go @@ -173,7 +173,11 @@ func TestWorkerRunOnceAdmitsOneQueuedInputAndLeavesBacklogRunnable(t *testing.T) secondID := createWorkerInput(t, ctx, store, agentID, userID, "second", now.Add(2*time.Second)) executor := newCaptureTurnExecutor(false) - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) worked, err := worker.RunOnce(ctx) if err != nil { t.Fatalf("worker run once: %v", err) @@ -223,7 +227,11 @@ func TestWorkerRunOnceBatchesSteeringInputsAndReleasesRuntime(t *testing.T) { ) executor := newCaptureTurnExecutor(false) - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) worked, err := worker.RunOnce(ctx) if err != nil { t.Fatalf("worker run once: %v", err) @@ -265,7 +273,11 @@ func TestWorkerRunOnceClaimsWakeupFromAnyProject(t *testing.T) { ) executor := newCaptureTurnExecutor(false) - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) worked, err := worker.RunOnce(ctx) if err != nil { t.Fatalf("worker run once: %v", err) @@ -292,7 +304,11 @@ func TestWorkerRunOnceDoesNotDoubleClaimActiveAgent(t *testing.T) { firstExecutor := newCaptureTurnExecutor(true) firstWorker := NewWorker( - store.Execution(), firstExecutor, Options{RuntimeLockLeaseDuration: 30 * time.Second, Capacity: 1}, + store.Execution(), firstExecutor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 30 * time.Second, + Capacity: 1, + }, ) firstCtx, cancelFirst := context.WithCancel(context.Background()) defer cancelFirst() @@ -318,7 +334,11 @@ func TestWorkerRunOnceDoesNotDoubleClaimActiveAgent(t *testing.T) { secondExecutor := newCaptureTurnExecutor(false) secondWorker := NewWorker( - store.Execution(), secondExecutor, Options{RuntimeLockLeaseDuration: 30 * time.Second, Capacity: 1}, + store.Execution(), secondExecutor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 30 * time.Second, + Capacity: 1, + }, ) worked, err := secondWorker.RunOnce(ctx) if err != nil { @@ -374,7 +394,12 @@ func TestWorkerCancelControlCancelsActiveTurnAndDeletesWakeup(t *testing.T) { executor := newCaptureTurnExecutor(true) worker := NewWorker( store.Execution(), executor, - Options{RuntimeLockLeaseDuration: 30 * time.Second, Capacity: 1, ControlSubscriber: bus}, + Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 30 * time.Second, + Capacity: 1, + ControlSubscriber: bus, + }, ) errs := make(chan error, 1) go func() { @@ -451,7 +476,11 @@ func TestWorkerCancelBeforeInitialRenewalCompletesCleanly(t *testing.T) { createWorkerInput(t, ctx, store, agentID, userID, "cancel before initial renewal", now.Add(time.Second)) executor := newCaptureTurnExecutor(false) - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 30 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 30 * time.Second, + Capacity: 1, + }) claim, worked, err := store.Execution().ClaimNextAgentWork( ctx, executionstore.ClaimNextAgentWorkInput{ @@ -538,7 +567,7 @@ func TestWorkerConcurrentRuntimesShareRouteAndCancelIndependently(t *testing.T) ) executor := newConcurrentTurnExecutor(2) - worker := NewWorker(store.Execution(), executor, Options{ + worker := NewWorker(store.Execution(), executor, Options{Continuation: &ContinuationOptions{}, RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 2, ControlSubscriber: bus, @@ -667,7 +696,12 @@ func TestWorkerRunContinuesAfterUnexpectedTurnCancellation(t *testing.T) { executor := &cancelOnceTurnExecutor{started: make(chan kernel.ModelWorkExecution, 2)} worker := NewWorker( store.Execution(), executor, - Options{RuntimeLockLeaseDuration: 30 * time.Second, Capacity: 1, ControlSubscriber: bus}, + Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 30 * time.Second, + Capacity: 1, + ControlSubscriber: bus, + }, ) runDone := make(chan error, 1) go func() { runDone <- worker.Run(ctx) }() @@ -709,7 +743,11 @@ func TestWorkerCancelFallsBackToRenewalWithoutControlDelivery(t *testing.T) { createWorkerInput(t, ctx, store, agentID, userID, "cancel by renewal", now.Add(time.Second)) executor := newCaptureTurnExecutor(true) - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) runDone := make(chan error, 1) go func() { worked, err := worker.RunOnce(ctx) @@ -787,7 +825,12 @@ func TestWorkerNewInputAfterCancelStartsNewTurn(t *testing.T) { blockingExecutor := newCaptureTurnExecutor(true) runningWorker := NewWorker( store.Execution(), blockingExecutor, - Options{RuntimeLockLeaseDuration: 30 * time.Second, Capacity: 1, ControlSubscriber: bus}, + Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 30 * time.Second, + Capacity: 1, + ControlSubscriber: bus, + }, ) errs := make(chan error, 1) go func() { @@ -838,7 +881,11 @@ func TestWorkerNewInputAfterCancelStartsNewTurn(t *testing.T) { ) nextExecutor := newCaptureTurnExecutor(false) nextWorker := NewWorker( - store.Execution(), nextExecutor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}, + store.Execution(), nextExecutor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }, ) worked, err := nextWorker.RunOnce(context.Background()) if err != nil { @@ -887,7 +934,11 @@ func TestWorkerRunOnceUsesRealKernelExecutor(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) worked, err := worker.RunOnce(ctx) if err != nil { t.Fatalf("worker real kernel run once: %v", err) @@ -942,7 +993,11 @@ func TestWorkerKernelStructuredQuestionBlocksResolvesAndContinues(t *testing.T) ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist question tool output") assertAgentInputAdmitted(t, ctx, pool, agentID, inputID) assertNoOpenInteractions(t, ctx, store, agentID) @@ -1026,7 +1081,11 @@ func TestWorkerKernelStructuredQuestionUsesCatalogAllowPolicy(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist catalog-allowed question output") assertAgentInputAdmitted(t, ctx, pool, agentID, inputID) requireWorkerClaim(t, ctx, worker, "dispatch catalog-allowed question") @@ -1106,7 +1165,11 @@ func TestWorkerKernelPermissionApprovalUnlocksStructuredQuestionTool(t *testing. ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist permission-gated question output") assertAgentInputAdmitted(t, ctx, pool, agentID, inputID) requireWorkerClaim(t, ctx, worker, "prepare question permission") @@ -1202,7 +1265,11 @@ func TestWorkerKernelPermissionDenialCompletesToolResult(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist permission-denied question output") requireWorkerClaim(t, ctx, worker, "prepare denied question permission") approval := assertOpenInteraction(t, ctx, store, agentID, "permission") @@ -1263,7 +1330,11 @@ func TestWorkerKernelMachineToolWithoutBindingFailsBeforeApproval(t *testing.T) ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist no-machine command output") requireWorkerClaim(t, ctx, worker, "fail no-machine command") requireWorkerClaim(t, ctx, worker, "continue after no-machine failure") @@ -1310,7 +1381,11 @@ func TestWorkerKernelDeniedMachineToolDoesNotRequireBinding(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist denied command output") requireWorkerClaim(t, ctx, worker, "deny command") requireWorkerClaim(t, ctx, worker, "continue after command denial") @@ -1363,7 +1438,11 @@ func TestWorkerKernelPermissionApprovalStartsMachineRunCommand(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist permission-gated command output") requireWorkerClaim(t, ctx, worker, "prepare command permission") approval := assertOpenInteraction(t, ctx, store, agentID, "permission") @@ -1432,7 +1511,11 @@ func TestWorkerKernelStructuredInteractionFormResolvesAtomically(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist multi-question output") requireWorkerClaim(t, ctx, worker, "dispatch multi-question tool") question := assertOpenInteraction(t, ctx, store, agentID, "question") @@ -1510,7 +1593,11 @@ func TestWorkerKernelMultipleToolCallsWaitForAllInteractions(t *testing.T) { ModelResolver: liveWorkerTestModelResolver(store, modelClient), Now: func() time.Time { return now.Add(2 * time.Second) }, } - worker := NewWorker(store.Execution(), executor, Options{RuntimeLockLeaseDuration: 15 * time.Second, Capacity: 1}) + worker := NewWorker(store.Execution(), executor, Options{ + Continuation: &ContinuationOptions{}, + RuntimeLockLeaseDuration: 15 * time.Second, + Capacity: 1, + }) requireWorkerClaim(t, ctx, worker, "persist sibling question outputs") requireWorkerClaim(t, ctx, worker, "dispatch sibling question tools") questions := assertOpenInteractions(t, ctx, store, agentID, "question", 2) diff --git a/internal/harness/worker/worker_test.go b/internal/harness/worker/worker_test.go index 3d185eed3..2050d148c 100644 --- a/internal/harness/worker/worker_test.go +++ b/internal/harness/worker/worker_test.go @@ -6,6 +6,7 @@ import ( "io" "log/slog" "strings" + "sync" "sync/atomic" "testing" "time" @@ -536,7 +537,7 @@ func TestWorkerLoopContinuesAfterTurnFailure(t *testing.T) { loopDone := make(chan struct{}) go func() { defer close(loopDone) - worker.runLoop(ctx) + runClaimLoop(ctx, worker) }() select { @@ -585,7 +586,7 @@ func TestWorkerLoopRetainsMixedShutdownError(t *testing.T) { Options{Log: logger, Capacity: 1}, ) - worker.runLoop(ctx) + runClaimLoop(ctx, worker) record := logs.String() if !strings.Contains(record, `"level":"error"`) { @@ -753,6 +754,7 @@ func TestRuntimeRenewalUsesLocalMonotonicBudgetAcrossClockSkew(t *testing.T) { uuid.UUID{1}, uuid.UUID{2}, executionstore.AgentRuntimeLockRecord{ID: uuid.UUID{3}, LeaseExpiresAt: test.leaseExpiresAt}, + time.Time{}, ) if err != nil { t.Fatalf("start runtime renewal: %v", err) @@ -793,6 +795,7 @@ func TestInitialRuntimeRenewalRetriesTransientFailure(t *testing.T) { uuid.UUID{1}, uuid.UUID{2}, executionstore.AgentRuntimeLockRecord{ID: uuid.UUID{3}, LeaseExpiresAt: leaseExpiresAt}, + time.Time{}, ) if err != nil { t.Fatalf("start runtime renewal after transient failure: %v", err) @@ -802,3 +805,179 @@ func TestInitialRuntimeRenewalRetriesTransientFailure(t *testing.T) { t.Fatalf("renewal calls = %d, want initial failure and successful retry", calls) } } + +func runClaimLoop(ctx context.Context, worker *Worker) { + var executions sync.WaitGroup + worker.claimLoop(ctx, make(chan struct{}, worker.capacity), &executions) + executions.Wait() +} + +type noopControlSubscriber struct{} + +func (noopControlSubscriber) SubscribeWorkerControl( + context.Context, + uuid.UUID, + func(context.Context, notifications.WorkerControl), +) (notifications.Subscription, error) { + return noopSubscription{}, nil +} + +type noopSubscription struct{} + +func (noopSubscription) Unsubscribe() error { return nil } + +type dispatchStore struct { + claimDelay time.Duration + activeClaims atomic.Int32 + maxClaims atomic.Int32 + claims atomic.Int32 +} + +func (s *dispatchStore) ClaimNextAgentWork( + ctx context.Context, + _ executionstore.ClaimNextAgentWorkInput, +) (executionstore.ClaimedAgentWork, bool, error) { + active := s.activeClaims.Add(1) + defer s.activeClaims.Add(-1) + raiseMax(&s.maxClaims, active) + select { + case <-time.After(s.claimDelay): + case <-ctx.Done(): + return executionstore.ClaimedAgentWork{}, false, ctx.Err() + } + s.claims.Add(1) + return executionstore.ClaimedAgentWork{ + ProjectID: uuid.New(), + AgentID: uuid.New(), + Kind: executionstore.AgentWorkModel, + RuntimeLock: executionstore.AgentRuntimeLockRecord{ID: uuid.New()}, + Model: executionstore.ClaimedModelWork{ + TurnID: uuid.New(), + InputIDs: []uuid.UUID{uuid.New()}, + OpeningEventSequence: 1, + }, + }, true, nil +} + +func (*dispatchStore) ReleaseAgentRuntimeLock(context.Context, uuid.UUID, uuid.UUID, uuid.UUID) error { + return nil +} + +func (*dispatchStore) AdvanceOwnedAgentWork( + context.Context, + executionstore.AdvanceOwnedAgentWorkInput, +) (executionstore.ClaimedAgentWork, bool, error) { + return executionstore.ClaimedAgentWork{}, false, nil +} + +func (*dispatchStore) RenewAgentRuntimeLock( + _ context.Context, + _, _, runtimeID uuid.UUID, + _ time.Duration, +) (executionstore.AgentRuntimeLockRenewal, error) { + return executionstore.AgentRuntimeLockRenewal{ + RuntimeLock: executionstore.AgentRuntimeLockRecord{ID: runtimeID}, + LocalLeaseBudgetStartedAt: time.Now(), + }, nil +} + +type holdingModelExecutor struct { + modelWorkOnlyExecutor + active atomic.Int32 + maxActive atomic.Int32 + release chan struct{} + started chan struct{} +} + +func (e *holdingModelExecutor) ExecuteModelWork(ctx context.Context, _ kernel.ModelWorkExecution) error { + raiseMax(&e.maxActive, e.active.Add(1)) + if e.started != nil { + e.started <- struct{}{} + } + defer e.active.Add(-1) + select { + case <-e.release: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func raiseMax(target *atomic.Int32, value int32) { + for { + current := target.Load() + if value <= current || target.CompareAndSwap(current, value) { + return + } + } +} + +func TestWorkerDispatcherBoundsClaimsAndNeverClaimsBeyondCapacity(t *testing.T) { + const capacity, claimConcurrency = 6, 2 + store := &dispatchStore{claimDelay: 20 * time.Millisecond} + executor := &holdingModelExecutor{release: make(chan struct{}), started: make(chan struct{}, capacity+1)} + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + worker := NewWorker(store, executor, Options{ + Log: log, + RuntimeLockLeaseDuration: time.Minute, + Capacity: capacity, + ClaimConcurrency: claimConcurrency, + ControlSubscriber: noopControlSubscriber{}, + }) + ctx, cancel := context.WithCancel(context.Background()) + runDone := make(chan error, 1) + go func() { runDone <- worker.Run(ctx) }() + + defer cancel() + for range capacity { + select { + case <-executor.started: + case <-time.After(2 * time.Second): + t.Fatal("dispatcher did not fill available slots") + } + } + select { + case <-executor.started: + t.Fatal("dispatcher started work without an available slot") + case <-time.After(100 * time.Millisecond): + } + if got := store.claims.Load(); got != capacity { + t.Fatalf("claims with every slot busy = %d, want %d", got, capacity) + } + executor.release <- struct{}{} + select { + case <-executor.started: + case <-time.After(2 * time.Second): + t.Fatal("dispatcher did not fill the released slot") + } + if got := store.claims.Load(); got != capacity+1 { + cancel() + t.Fatalf("claims after one slot freed = %d, want %d", got, capacity+1) + } + cancel() + if err := <-runDone; !errors.Is(err, context.Canceled) { + t.Fatalf("run error = %v, want context canceled", err) + } + if got := store.maxClaims.Load(); got > claimConcurrency { + t.Fatalf("max concurrent claims = %d, want at most %d", got, claimConcurrency) + } + if got := executor.maxActive.Load(); got > capacity { + t.Fatalf("max concurrent executions = %d, want at most %d", got, capacity) + } +} + +func (s *retainedRuntimeStore) AdvanceOwnedAgentWork( + ctx context.Context, + input executionstore.AdvanceOwnedAgentWorkInput, +) (executionstore.ClaimedAgentWork, bool, error) { + err := s.ReleaseAgentRuntimeLock(ctx, input.ProjectID, input.AgentID, input.RuntimeLockID) + return executionstore.ClaimedAgentWork{}, false, err +} + +func (s *renewalDeadlineStore) AdvanceOwnedAgentWork( + ctx context.Context, + input executionstore.AdvanceOwnedAgentWorkInput, +) (executionstore.ClaimedAgentWork, bool, error) { + err := s.ReleaseAgentRuntimeLock(ctx, input.ProjectID, input.AgentID, input.RuntimeLockID) + return executionstore.ClaimedAgentWork{}, false, err +} diff --git a/internal/httpapi/agent_interactions_integration_test.go b/internal/httpapi/agent_interactions_integration_test.go index c0c7be955..bfc7699f0 100644 --- a/internal/httpapi/agent_interactions_integration_test.go +++ b/internal/httpapi/agent_interactions_integration_test.go @@ -163,12 +163,10 @@ func TestPublicAgentInteractionResolveMarksWakeup(t *testing.T) { ).Scan(&wakeups); err != nil { t.Fatalf("query wakeup: %v", err) } - if wakeups != 1 { - t.Fatalf( - "interaction resolution should mark exactly one agent wakeup, got %d", - wakeups, - ) + if wakeups != 0 { + t.Fatalf("owned interaction resolution published %d wakeups", wakeups) } + assertResolvedInteractionReleasedWakeup(t, ctx, pool, store, project.ProjectUUID, agentID) assertInteractionResponseAgentInput(t, ctx, pool, project.ProjectUUID, agentID, interactionID) assertInteractionResponseLedgerEvent(t, ctx, pool, project.ProjectUUID, agentID, interactionID) @@ -1540,12 +1538,10 @@ GROUP BY tool_call.state, result.outcome ).Scan(&wakeups); err != nil { t.Fatalf("query denial wakeup: %v", err) } - if wakeups != 1 { - t.Fatalf( - "permission denial should mark one wakeup after storage-owned tool-result admission, got %d", - wakeups, - ) + if wakeups != 0 { + t.Fatalf("owned permission denial published %d wakeups", wakeups) } + assertResolvedInteractionReleasedWakeup(t, ctx, pool, store, project.ProjectUUID, agentID) } func TestPublicAgentInteractionRejectsUnknownPermissionOption( @@ -2463,3 +2459,33 @@ func createHTTPCustomToolCallForAgent( } return toolCall.ID } + +func assertResolvedInteractionReleasedWakeup( + t *testing.T, ctx context.Context, pool *pgxpool.Pool, store *storage.Store, projectID, agentID uuid.UUID, +) { + t.Helper() + var runtimeID uuid.UUID + var ready bool + if err := pool.QueryRow(ctx, ` +SELECT lock.id, head.logical_ready_at IS NOT NULL +FROM agent_runtime_locks lock +JOIN agent_execution_state head ON head.agent_id = lock.agent_id +WHERE lock.agent_id = $1`, agentID).Scan(&runtimeID, &ready); err != nil { + t.Fatalf("read owned interaction execution: %v", err) + } + if !ready { + t.Fatal("resolved interaction has no scheduling obligation") + } + if err := store.Execution().ReleaseAgentRuntimeLock(ctx, projectID, agentID, runtimeID); err != nil { + t.Fatalf("release interaction runtime: %v", err) + } + var wakeups int + if err := pool.QueryRow( + ctx, `SELECT count(*) FROM agent_wakeups WHERE agent_id = $1`, agentID, + ).Scan(&wakeups); err != nil { + t.Fatalf("read released interaction wakeup: %v", err) + } + if wakeups != 1 { + t.Fatalf("released interaction wakeups=%d, want 1", wakeups) + } +} diff --git a/internal/httpapi/agents_list_route_integration_test.go b/internal/httpapi/agents_list_route_integration_test.go index ba36c7922..b6a2cb62c 100644 --- a/internal/httpapi/agents_list_route_integration_test.go +++ b/internal/httpapi/agents_list_route_integration_test.go @@ -503,15 +503,20 @@ func seedListAgentAt( createdAt time.Time, ) executionstore.AgentRecord { t.Helper() - var id uuid.UUID + id := uuid.New() if err := pool.QueryRow(ctx, ` -INSERT INTO agents ( - org_id, project_id, state, name, current_config_id, - idempotency_key, created_at, updated_at +WITH inserted AS ( + INSERT INTO agents ( + id, root_agent_id, org_id, project_id, state, name, current_config_id, + idempotency_key, created_at, updated_at + ) + VALUES ($7, $7, $1, $2, 'active', $3, $4, $5, $6, $6) + RETURNING id ) -VALUES ($1, $2, 'active', $3, $4, $5, $6, $6) -RETURNING id -`, project.OrgUUID, project.ProjectUUID, name, configID, idempotencyKey, createdAt).Scan(&id); err != nil { +INSERT INTO agent_execution_state (agent_id, turn_continuable, incomplete_tools) +SELECT id, false, false FROM inserted +RETURNING agent_id +`, project.OrgUUID, project.ProjectUUID, name, configID, idempotencyKey, createdAt, id).Scan(&id); err != nil { t.Fatalf("seed agent at creation time: %v", err) } record, err := store.Execution().GetAgentInProject(ctx, project.ProjectUUID, id) diff --git a/internal/httpapi/discord_app_setup.go b/internal/httpapi/discord_app_setup.go deleted file mode 100644 index 751f35b18..000000000 --- a/internal/httpapi/discord_app_setup.go +++ /dev/null @@ -1,7 +0,0 @@ -package httpapi - -import "github.com/omnara-ai/omnara/internal/integration/discord" - -func WithDiscordClientConfig(config discord.Config) Option { - return func(s *Server) { s.discordClientConfig = config } -} diff --git a/internal/httpapi/github_app_setup.go b/internal/httpapi/github_app_setup.go deleted file mode 100644 index 35e01e2ab..000000000 --- a/internal/httpapi/github_app_setup.go +++ /dev/null @@ -1,7 +0,0 @@ -package httpapi - -import "github.com/omnara-ai/omnara/internal/integration/github" - -func WithGitHubClientConfig(config github.Config) Option { - return func(s *Server) { s.githubClientConfig = config } -} diff --git a/internal/httpapi/model_context_test_helpers_test.go b/internal/httpapi/model_context_test_helpers_test.go index 66f0664d4..42ba043c5 100644 --- a/internal/httpapi/model_context_test_helpers_test.go +++ b/internal/httpapi/model_context_test_helpers_test.go @@ -23,14 +23,15 @@ func claimNormalModelCallForHTTPTest( inputEventSequence int64, ) executionstore.ModelCallClaim { t.Helper() - claim, err := store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtime.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agentConfigID, - InputEventSequence: inputEventSequence, - }) + prepared1, err := store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: projectID, + AgentID: agentID, + RuntimeLockID: runtime.ID, + OpeningInputIDs: openingInputIDs, + }) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim normal model call: %v", err) } diff --git a/internal/httpapi/options_integration_test.go b/internal/httpapi/options_integration_test.go new file mode 100644 index 000000000..e1af98306 --- /dev/null +++ b/internal/httpapi/options_integration_test.go @@ -0,0 +1,23 @@ +//go:build integration + +package httpapi + +import ( + "github.com/benbjohnson/clock" + "github.com/omnara-ai/omnara/internal/integration/discord" + "github.com/omnara-ai/omnara/internal/integration/github" +) + +func WithTimer(timer clock.Clock) Option { + return func(s *Server) { + s.timer = timer + } +} + +func WithDiscordClientConfig(config discord.Config) Option { + return func(s *Server) { s.discordClientConfig = config } +} + +func WithGitHubClientConfig(config github.Config) Option { + return func(s *Server) { s.githubClientConfig = config } +} diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index aeaf1db47..27d96c407 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -94,12 +94,6 @@ type Server struct { timer clock.Clock } -func WithTimer(timer clock.Clock) Option { - return func(s *Server) { - s.timer = timer - } -} - type daemonNotificationConfig struct { subscriber notifications.DaemonWakeupSubscriber presence notifications.DaemonPresenceStore @@ -444,13 +438,19 @@ func New(log *slog.Logger, store *storage.Store, opts ...Option) (*Server, error ClientMetadataHTTPClient: server.oauthClientMetadataHTTPClient, }) if server.agentEventWakeupSubscriber == nil { - return nil, fmt.Errorf("agent event wakeup subscriber is required; wire via WithAgentEventWakeupSubscriber") + return nil, fmt.Errorf( + "agent event wakeup subscriber is required; wire via WithAgentEventWakeupSubscriber", + ) } if server.agentToolCallUpdateSubscriber == nil { - return nil, fmt.Errorf("agent tool call update subscriber is required; wire via WithAgentToolCallUpdateSubscriber") + return nil, fmt.Errorf( + "agent tool call update subscriber is required; wire via WithAgentToolCallUpdateSubscriber", + ) } if server.agentStreamDeltaSubscriber == nil { - return nil, fmt.Errorf("agent stream delta subscriber is required; wire via WithAgentStreamDeltaSubscriber") + return nil, fmt.Errorf( + "agent stream delta subscriber is required; wire via WithAgentStreamDeltaSubscriber", + ) } if store != nil { server.agentEventStreamReconciler = newAgentEventStreamReconciler( diff --git a/internal/httpapi/slack_event_routes_integration_test.go b/internal/httpapi/slack_event_routes_integration_test.go index e4b8077a7..370e35db5 100644 --- a/internal/httpapi/slack_event_routes_integration_test.go +++ b/internal/httpapi/slack_event_routes_integration_test.go @@ -1093,7 +1093,11 @@ func testSlackEventsOpenInteractionContinuesWithNewMessage(t *testing.T, kind st var stopEvents int if err := pool.QueryRow( ctx, - `SELECT count(*)::int FROM agent_stop_events WHERE project_id = $1 AND agent_id = $2`, + `SELECT count(*)::int FROM agent_events event +JOIN agent_inputs input ON input.agent_id = event.agent_id AND input.id = event.agent_input_id +JOIN agents agent ON agent.id = event.agent_id +WHERE agent.project_id = $1 AND event.agent_id = $2 AND input.input_kind = 'control' AND +input.control_type = 'cancel_current'`, fixture.Project.ProjectUUID, integrationTarget.AgentID, ).Scan(&stopEvents); err != nil { diff --git a/internal/interactionform/answers.go b/internal/interactionform/answers.go new file mode 100644 index 000000000..44a44975a --- /dev/null +++ b/internal/interactionform/answers.go @@ -0,0 +1,47 @@ +package interactionform + +type RenderedAnswers struct { + Answers []RenderedAnswer `json:"answers"` +} + +type RenderedAnswer struct { + QuestionIndex int `json:"question_index"` + Question string `json:"question"` + SelectedOptions []SelectedAnswerOption `json:"selected_options"` + Text string `json:"text,omitempty"` +} + +type SelectedAnswerOption struct { + OptionIndex int `json:"option_index"` + Label string `json:"label"` +} + +func RenderAnswers( + form Form, + resolution Resolution, +) RenderedAnswers { + result := RenderedAnswers{ + Answers: make([]RenderedAnswer, 0, len(resolution.Answers)), + } + for questionIndex, answer := range resolution.Answers { + question := form.Questions[questionIndex] + selectedOptions := make( + []SelectedAnswerOption, + 0, + len(answer.OptionIndices), + ) + for _, optionIndex := range answer.OptionIndices { + selectedOptions = append(selectedOptions, SelectedAnswerOption{ + OptionIndex: optionIndex, + Label: question.Options[optionIndex].Label, + }) + } + result.Answers = append(result.Answers, RenderedAnswer{ + QuestionIndex: questionIndex, + Question: question.Prompt, + SelectedOptions: selectedOptions, + Text: answer.Text, + }) + } + return result +} diff --git a/internal/interactionform/form_test.go b/internal/interactionform/form_test.go index 583271da6..3fafc6528 100644 --- a/internal/interactionform/form_test.go +++ b/internal/interactionform/form_test.go @@ -178,3 +178,30 @@ func resolutionTestForm(t *testing.T) Form { } return value } + +func TestQuestionToolResultKeepsPositionalIdentityForDuplicateLabels(t *testing.T) { + result := RenderAnswers( + Form{Questions: []Question{{ + Prompt: "Choose", + Multiple: true, + Options: []Option{ + {Label: "Other"}, + {Label: "Other", AllowsText: true}, + }, + }}}, + Resolution{Answers: []Answer{{ + OptionIndices: []int{0, 1}, + }}}, + ) + if len(result.Answers) != 1 || result.Answers[0].QuestionIndex != 0 { + t.Fatalf("question result = %+v", result) + } + selected := result.Answers[0].SelectedOptions + if len(selected) != 2 || + selected[0].OptionIndex != 0 || + selected[1].OptionIndex != 1 || + selected[0].Label != "Other" || + selected[1].Label != "Other" { + t.Fatalf("selected options = %+v", selected) + } +} diff --git a/internal/interactionform/question.go b/internal/interactionform/question.go new file mode 100644 index 000000000..1fabe08a0 --- /dev/null +++ b/internal/interactionform/question.go @@ -0,0 +1,26 @@ +package interactionform + +import ( + "fmt" + "strings" +) + +func RenderQuestion(form Form) string { + var builder strings.Builder + builder.WriteString("Question: ") + builder.WriteString(form.Title) + for _, item := range form.Context { + builder.WriteString("\n") + builder.WriteString(item.Label) + builder.WriteString(": ") + builder.WriteString(item.Value) + } + for index, question := range form.Questions { + builder.WriteString(fmt.Sprintf("\n%d. %s", index+1, question.Prompt)) + for _, option := range question.Options { + builder.WriteString("\n - ") + builder.WriteString(option.Label) + } + } + return builder.String() +} diff --git a/internal/machinepool/manager_integration_test.go b/internal/machinepool/manager_integration_test.go index b6a402ac3..94004126c 100644 --- a/internal/machinepool/manager_integration_test.go +++ b/internal/machinepool/manager_integration_test.go @@ -232,26 +232,28 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi } machineSecretID := machineSecret.ID.String() maxCPU, maxMemoryMB := 100, 1024*1024 - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Provider Config Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{"mode":"provision"}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 4, - MaxTotalCPU: new(maxCPU), - MaxTotalMemoryMB: new(maxMemoryMB), - MaxMachineCPU: new(maxCPU), - MaxMachineMemoryMB: new(1024), - }, - 1, - 1024, - map[string]string{"PLAIN": "plain"}, - map[string]string{"API_TOKEN": machineSecretID}, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Provider Config Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{"mode":"provision"}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 4, + MaxTotalCPU: new(maxCPU), + MaxTotalMemoryMB: new(maxMemoryMB), + MaxMachineCPU: new(maxCPU), + MaxMachineMemoryMB: new(1024), + }, + 1, + 1024, + map[string]string{"PLAIN": "plain"}, + map[string]string{"API_TOKEN": machineSecretID}, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -304,9 +306,12 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi } provider := &captureProvider{ provisionResourceID: "resource-1", - prepare: func(provisioning executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) { + prepare: func(provisioning executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) { if provisioning.CPU != nil || provisioning.MemoryMB != nil { - t.Fatalf("preparation input resources = cpu %v memory %v, want unresolved", provisioning.CPU, provisioning.MemoryMB) + t.Fatalf("preparation input resources = cpu %v memory %v, want unresolved", + provisioning.CPU, + provisioning.MemoryMB) } return executionstore.MachineResourceFacts{ CPU: new(1), @@ -399,7 +404,8 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi ); err != nil { t.Fatalf("insert second project machine grant: %v", err) } - provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) { + provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) { return executionstore.MachineResourceFacts{ CPU: new(1), MemoryMB: new(2048), @@ -448,7 +454,8 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi ); err != nil { t.Fatalf("insert drift project machine grant: %v", err) } - provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) { + provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) { return executionstore.MachineResourceFacts{ CPU: new(1), MemoryMB: new(1024), @@ -459,7 +466,8 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi if err := manager.ProvisionMachine(ctx, orgID, driftMachineID); !errors.Is(err, provisionErr) { t.Fatalf("initial provider failure = %v, want provision failure", err) } - provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) { + provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) { return executionstore.MachineResourceFacts{ CPU: new(2), MemoryMB: new(1024), @@ -468,7 +476,10 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi provider.provisionErr = nil provider.provisioning = nil makePoolMachineReadyForManagerReconcile(t, ctx, pool, orgID, driftMachineID) - if err := manager.ProvisionMachine(ctx, orgID, driftMachineID); !errors.Is(err, storeerr.ErrStateTransitionConflict) { + if err := manager.ProvisionMachine(ctx, + orgID, + driftMachineID); !errors.Is(err, + storeerr.ErrStateTransitionConflict) { t.Fatalf("provider fact drift error = %v, want state transition conflict", err) } driftMachine, err := store.Execution().GetMachine(ctx, orgID, driftMachineID) @@ -487,7 +498,10 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi provider.inspectResourceID = "resource-from-ambiguous-first-attempt" provider.inspectFound = true makePoolMachineReadyForManagerReconcile(t, ctx, pool, orgID, driftMachineID) - if err := manager.ProvisionMachine(ctx, orgID, driftMachineID); !errors.Is(err, storeerr.ErrStateTransitionConflict) { + if err := manager.ProvisionMachine(ctx, + orgID, + driftMachineID); !errors.Is(err, + storeerr.ErrStateTransitionConflict) { t.Fatalf("final provider fact drift error = %v, want state transition conflict", err) } driftMachine, err = store.Execution().GetMachine(ctx, orgID, driftMachineID) @@ -538,7 +552,8 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi ); err != nil { t.Fatalf("seed final provisioning attempt: %v", err) } - provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) { + provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) { return executionstore.MachineResourceFacts{ CPU: new(1), MemoryMB: new(1024), @@ -595,7 +610,10 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi provider.provisionResourceID = "resource-with-permanent-error" provider.provisionErr = permanentErr provider.deletedResourceIDs = nil - if err := manager.ProvisionMachine(ctx, orgID, permanentMachineID); !errors.Is(err, providers.ErrPermanent) { + if err := manager.ProvisionMachine(ctx, + orgID, + permanentMachineID); !errors.Is(err, + providers.ErrPermanent) { t.Fatalf("permanent provider failure = %v, want permanent error", err) } permanentMachine, err := store.Execution().GetMachine(ctx, orgID, permanentMachineID) @@ -620,13 +638,17 @@ func testPoolMachineManagerProvisioningScenario(t *testing.T, scenario poolMachi ); err != nil { t.Fatalf("seed final admission attempt: %v", err) } - provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) { + provider.prepare = func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) { return executionstore.MachineResourceFacts{ CPU: new(1), MemoryMB: new(2048), }, nil } - if err := manager.ProvisionMachine(ctx, orgID, machineID); !errors.Is(err, storeerr.ErrStateTransitionConflict) { + if err := manager.ProvisionMachine(ctx, + orgID, + machineID); !errors.Is(err, + storeerr.ErrStateTransitionConflict) { t.Fatalf("final admission error = %v, want state transition conflict", err) } machine, err := store.Execution().GetMachine(ctx, orgID, machineID) @@ -683,26 +705,28 @@ func TestManagerValidatesPoolPolicyBeforeProvisioning(t *testing.T) { "provider-policy-auth", "pool-token", ) - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Provider Policy Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(100), - MaxTotalMemoryMB: new(1024 * 1024), - MaxMachineCPU: new(100), - MaxMachineMemoryMB: new(1024 * 1024), - }, - 1, - 1024, - nil, - nil, - map[string]any{"image": "default"}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Provider Policy Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(100), + MaxTotalMemoryMB: new(1024 * 1024), + MaxMachineCPU: new(100), + MaxMachineMemoryMB: new(1024 * 1024), + }, + 1, + 1024, + nil, + nil, + map[string]any{"image": "default"}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) } @@ -805,26 +829,28 @@ func TestManagerProvisioningEnvironmentSecrets(t *testing.T) { t.Fatalf("create machine secret: %v", err) } machineSecretID := machineSecret.ID.String() - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Secret Resolution Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{"mode":"provision"}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(100), - MaxTotalMemoryMB: new(1024 * 1024), - MaxMachineCPU: new(100), - MaxMachineMemoryMB: new(1024 * 1024), - }, - 1, - 1024, - map[string]string{"PLAIN": "plain"}, - map[string]string{"API_TOKEN": machineSecretID}, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Secret Resolution Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{"mode":"provision"}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(100), + MaxTotalMemoryMB: new(1024 * 1024), + MaxMachineCPU: new(100), + MaxMachineMemoryMB: new(1024 * 1024), + }, + 1, + 1024, + map[string]string{"PLAIN": "plain"}, + map[string]string{"API_TOKEN": machineSecretID}, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -892,7 +918,10 @@ func TestManagerProvisioningEnvironmentSecrets(t *testing.T) { err = manager.ProvisionMachine(ctx, orgID, machineID) if scenario == "invalid_value" { - if !errors.Is(err, storeerr.ErrPermanentEnvironment) || !strings.Contains(err.Error(), "secret_env.API_TOKEN") { + if !errors.Is(err, + storeerr.ErrPermanentEnvironment) || + !strings.Contains(err.Error(), + "secret_env.API_TOKEN") { t.Fatalf("provision machine error = %v, want permanent environment failure", err) } } else if err != nil { @@ -945,26 +974,28 @@ func TestManagerUsesArchivedPoolProviderConfigForCleanup(t *testing.T) { "cleanup-token", ) maxCPU, maxMemoryMB := 100, 1024*1024 - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Cleanup Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{"mode":"cleanup"}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(maxCPU), - MaxTotalMemoryMB: new(maxMemoryMB), - MaxMachineCPU: new(maxCPU), - MaxMachineMemoryMB: new(maxMemoryMB), - }, - 1, - 1024, - map[string]string{}, - nil, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Cleanup Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{"mode":"cleanup"}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(maxCPU), + MaxTotalMemoryMB: new(maxMemoryMB), + MaxMachineCPU: new(maxCPU), + MaxMachineMemoryMB: new(maxMemoryMB), + }, + 1, + 1024, + map[string]string{}, + nil, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -1334,26 +1365,28 @@ func TestManagerDeletesMachineWithoutProviderProvisionAttempt(t *testing.T) { "must-not-load", ) maxCPU, maxMemoryMB := 100, 1024*1024 - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Unattempted Deleting Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{"api_token":"must-not-load"}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(maxCPU), - MaxTotalMemoryMB: new(maxMemoryMB), - MaxMachineCPU: new(maxCPU), - MaxMachineMemoryMB: new(maxMemoryMB), - }, - 1, - 1024, - map[string]string{}, - nil, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Unattempted Deleting Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{"api_token":"must-not-load"}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(maxCPU), + MaxTotalMemoryMB: new(maxMemoryMB), + MaxMachineCPU: new(maxCPU), + MaxMachineMemoryMB: new(maxMemoryMB), + }, + 1, + 1024, + map[string]string{}, + nil, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -1472,26 +1505,28 @@ func TestManagerRetriesDeletingAttemptedMachineWithoutResource(t *testing.T) { "attempted-delete-token", ) maxCPU, maxMemoryMB := 100, 1024*1024 - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Attempted Cleanup Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(maxCPU), - MaxTotalMemoryMB: new(maxMemoryMB), - MaxMachineCPU: new(maxCPU), - MaxMachineMemoryMB: new(maxMemoryMB), - }, - 1, - 1024, - map[string]string{}, - nil, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Attempted Cleanup Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(maxCPU), + MaxTotalMemoryMB: new(maxMemoryMB), + MaxMachineCPU: new(maxCPU), + MaxMachineMemoryMB: new(maxMemoryMB), + }, + 1, + 1024, + map[string]string{}, + nil, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -1546,7 +1581,8 @@ func TestManagerRetriesDeletingAttemptedMachineWithoutResource(t *testing.T) { if err != nil { t.Fatalf("get delete failed machine: %v", err) } - if machine.LifecycleState != "delete_failed" || machine.LifecycleReasonCode != "provider_resource_not_found" || + if machine.LifecycleState != "delete_failed" || + machine.LifecycleReasonCode != "provider_resource_not_found" || machine.DeleteAttempts != 1 { t.Fatalf("machine did not enter provider missing-resource retry: %+v", machine) } @@ -1576,26 +1612,28 @@ func TestManagerFinalizesStaleMissingProviderResource(t *testing.T) { "stale-missing-token", ) maxCPU, maxMemoryMB := 100, 1024*1024 - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Stale Missing Provider Resource Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(maxCPU), - MaxTotalMemoryMB: new(maxMemoryMB), - MaxMachineCPU: new(maxCPU), - MaxMachineMemoryMB: new(maxMemoryMB), - }, - 1, - 1024, - map[string]string{}, - nil, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Stale Missing Provider Resource Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(maxCPU), + MaxTotalMemoryMB: new(maxMemoryMB), + MaxMachineCPU: new(maxCPU), + MaxMachineMemoryMB: new(maxMemoryMB), + }, + 1, + 1024, + map[string]string{}, + nil, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -1674,26 +1712,28 @@ func TestManagerReconcilesDeletingMachineWithProviderResource(t *testing.T) { "deleting-resource-token", ) maxCPU, maxMemoryMB := 100, 1024*1024 - machinePool, err := store.Execution().CreateMachinePool(ctx, machinePoolInputWithDefaultMachineForManagerTest( - t, - executionstore.CreateMachinePoolInput{ - OrgID: orgID, - Name: "Deleting Resource Pool", - Provider: "capture", - ProviderConfig: json.RawMessage(`{"api_token":"deleting-token"}`), - ProviderAuthSecretID: providerAuthSecretID, - MaxTotalMachines: 1, - MaxTotalCPU: new(maxCPU), - MaxTotalMemoryMB: new(maxMemoryMB), - MaxMachineCPU: new(maxCPU), - MaxMachineMemoryMB: new(maxMemoryMB), - }, - 1, - 1024, - map[string]string{}, - nil, - map[string]any{}, - )) + machinePool, + err := store.Execution().CreateMachinePool(ctx, + machinePoolInputWithDefaultMachineForManagerTest( + t, + executionstore.CreateMachinePoolInput{ + OrgID: orgID, + Name: "Deleting Resource Pool", + Provider: "capture", + ProviderConfig: json.RawMessage(`{"api_token":"deleting-token"}`), + ProviderAuthSecretID: providerAuthSecretID, + MaxTotalMachines: 1, + MaxTotalCPU: new(maxCPU), + MaxTotalMemoryMB: new(maxMemoryMB), + MaxMachineCPU: new(maxCPU), + MaxMachineMemoryMB: new(maxMemoryMB), + }, + 1, + 1024, + map[string]string{}, + nil, + map[string]any{}, + )) if err != nil { t.Fatalf("create machine pool: %v", err) @@ -1829,7 +1869,10 @@ func seedManagerProjectActor( } if _, err := store.Identity().AddProjectMembership( ctx, - identitystore.AddProjectMembershipInput{OrgID: orgID, ProjectID: projectID, UserID: actor.ID, Role: "admin"}, + identitystore.AddProjectMembershipInput{OrgID: orgID, + ProjectID: projectID, + UserID: actor.ID, + Role: "admin"}, ); err != nil { t.Fatalf("add actor project membership: %v", err) } @@ -2146,8 +2189,10 @@ VALUES ($1, $2, $3, $4, `, configID, orgID, projectID, configuredModelID, now) exec("insert cleanup agent", ` -INSERT INTO agents(id, org_id, project_id, state, name, current_config_id, created_at, updated_at) -VALUES ($1, $2, $3, 'active', 'Manager Cleanup Agent', $4, $5, $5) +WITH inserted AS (INSERT INTO agents(id, root_agent_id, org_id, project_id, state, name, +current_config_id, created_at, updated_at) +VALUES ($1, $1, $2, $3, 'active', 'Manager Cleanup Agent', $4, $5, $5) RETURNING id) +INSERT INTO agent_execution_state(agent_id,turn_continuable,incomplete_tools) SELECT id,false,false FROM inserted `, agentID, orgID, projectID, configID, now) exec("insert cleanup pool grant", ` @@ -2260,14 +2305,15 @@ func testProviderCatalog(definition providers.Definition) Catalog { } type captureProvider struct { - provisionResourceID string - provisioning *executionstore.MachineProvisioningConfig - machineEnv map[string]string - machineToken string - firstAttempt bool - provisionErr error - validateConfigErr error - prepare func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, error) + provisionResourceID string + provisioning *executionstore.MachineProvisioningConfig + machineEnv map[string]string + machineToken string + firstAttempt bool + provisionErr error + validateConfigErr error + prepare func(executionstore.MachineProvisioningConfig) (executionstore.MachineResourceFacts, + error) installationID uuid.UUID machineID uuid.UUID inspectResourceID string @@ -2370,7 +2416,11 @@ func (p *captureProvider) DeleteMachine( } // Only this test database contains the trigger and its nontransactional counter. -func installWakeCommitFailure(t *testing.T, ctx context.Context, pool *pgxpool.Pool, code string, failures int) { +func installWakeCommitFailure(t *testing.T, + ctx context.Context, + pool *pgxpool.Pool, + code string, + failures int) { t.Helper() _, err := pool.Exec(ctx, fmt.Sprintf(` CREATE SEQUENCE test_wake_attempts; diff --git a/internal/modelcontext/context.go b/internal/modelcontext/context.go index 681565c5b..94b6b029d 100644 --- a/internal/modelcontext/context.go +++ b/internal/modelcontext/context.go @@ -44,6 +44,13 @@ func (b Builder) Build(ctx context.Context, input BuildInput) (Bundle, error) { snapshot = captured } watermark := snapshot.InputEventSequence + if input.ContextData != nil && input.CheckpointOverride == nil { + data := input.ContextData + if data.ProjectID != input.ProjectID || data.AgentID != input.AgentID || data.InputEventSequence != watermark { + return Bundle{}, fmt.Errorf("prepared model context does not match the requested frontier") + } + b.Store = preparedContextStore{Store: b.Store, data: data} + } afterSequence := int64(0) var checkpointRef *CheckpointRef if input.CheckpointOverride != nil { diff --git a/internal/modelcontext/prepared.go b/internal/modelcontext/prepared.go new file mode 100644 index 000000000..0c1c201b9 --- /dev/null +++ b/internal/modelcontext/prepared.go @@ -0,0 +1,52 @@ +package modelcontext + +import ( + "context" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/executionstore" +) + +type preparedContextStore struct { + Store + data *executionstore.ModelContextData +} + +func (s preparedContextStore) GetLatestApplicableContextCheckpoint( + _ context.Context, _, _ uuid.UUID, _ int64, +) (executionstore.ContextCheckpointRecord, bool, error) { + if s.data.Checkpoint == nil { + return executionstore.ContextCheckpointRecord{}, false, nil + } + return *s.data.Checkpoint, true, nil +} + +func (s preparedContextStore) IsOutputLimitBoundary( + ctx context.Context, projectID, agentID uuid.UUID, sequence int64, +) (bool, error) { + if s.data.Checkpoint != nil && sequence == s.data.Checkpoint.SummarizedThroughEventSequence { + return s.data.CheckpointEndsWithOutputLimit, nil + } + return s.Store.IsOutputLimitBoundary(ctx, projectID, agentID, sequence) +} + +func (s preparedContextStore) ListContextEvents( + _ context.Context, _, _ uuid.UUID, after, watermark int64, limit int32, +) ([]executionstore.ContextEventRecord, error) { + var result []executionstore.ContextEventRecord + for _, event := range s.data.Events { + if event.Sequence > after && event.Sequence <= watermark { + result = append(result, event) + if len(result) == int(limit) { + break + } + } + } + return result, nil +} + +func (s preparedContextStore) ListCompletedToolCallsAtWatermark( + _ context.Context, _, _ uuid.UUID, _, _ int64, +) ([]executionstore.ToolCallRecord, error) { + return s.data.ToolCalls, nil +} diff --git a/internal/modelcontext/prepared_test.go b/internal/modelcontext/prepared_test.go new file mode 100644 index 000000000..6b4452ae5 --- /dev/null +++ b/internal/modelcontext/prepared_test.go @@ -0,0 +1,69 @@ +package modelcontext + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/stretchr/testify/require" +) + +func TestPreparedContextMatchesStoredTranscriptAndCheckpointOverride(t *testing.T) { + ctx := t.Context() + for _, checkpoint := range []bool{false, true} { + name := "full transcript" + if checkpoint { + name = "checkpoint" + } + t.Run(name, func(t *testing.T) { + store := &fakeContextStore{watermark: 602} + for i := 1; i <= 601; i++ { + store.messages = append(store.messages, contextTextEvent(t, testIDN(i), int64(i), "message")) + } + after := int64(0) + if checkpoint { + after = 42 + store.checkpoints = []executionstore.ContextCheckpointRecord{{ID: testIDN(700), Summary: "Earlier work", + SummarizedThroughEventSequence: after, CheckpointEventSequence: 43}} + store.outputLimitBoundaries = map[int64]bool{after: true} + } + store.toolCalls = []executionstore.ToolCallRecord{{ + ID: testIDN(800), ProviderCallID: "tool", TurnID: testTurnID, ModelCallContextID: testIDN(801), + ToolCallResultID: testIDN(802), ToolResultEventID: testIDN(803), Name: "list_agents", Input: []byte(`{}`), + SourceEventSequence: 601, ToolResultEventSequence: 602, State: executionstore.ToolCallStateCompleted, + Outcome: executionstore.ToolResultOutcomeSucceeded, ResultContentParts: []byte(`[{"type":"text","text":"done"}]`), + }} + snapshot, err := store.CaptureAgentConfigForModelContext(ctx, testProjectID, testAgentID) + require.NoError(t, err) + data := &executionstore.ModelContextData{ProjectID: testProjectID, AgentID: testAgentID, + InputEventSequence: 602, Events: store.messages[int(after):], ToolCalls: store.toolCalls} + if checkpoint { + data.Checkpoint = &store.checkpoints[0] + data.CheckpointEndsWithOutputLimit = true + } + input := BuildInput{ProjectID: testProjectID, AgentID: testAgentID, TurnID: testTurnID, + OpeningInputIDs: []uuid.UUID{testInputID}, Now: time.Now(), AgentConfigSnapshot: &snapshot} + builder := Builder{Store: store} + expected, err := builder.Build(ctx, input) + require.NoError(t, err) + input.ContextData = data + prepared, err := builder.Build(ctx, input) + require.NoError(t, err) + require.Equal(t, expected, prepared) + require.Len(t, prepared.Messages, 601-int(after)) + input.CheckpointOverride = &CheckpointRef{Summary: "Candidate", SummarizedThroughEventSequence: 500} + prepared, err = builder.Build(ctx, input) + require.NoError(t, err) + input.ContextData = nil + expected, err = builder.Build(ctx, input) + require.NoError(t, err) + require.Equal(t, expected, prepared) + input.ContextData = data + input.CheckpointOverride = nil + data.InputEventSequence++ + _, err = builder.Build(ctx, input) + require.ErrorContains(t, err, "does not match") + }) + } +} diff --git a/internal/modelcontext/types.go b/internal/modelcontext/types.go index d561a790d..4fd195fc5 100644 --- a/internal/modelcontext/types.go +++ b/internal/modelcontext/types.go @@ -18,6 +18,7 @@ type BuildInput struct { OpeningInputIDs []uuid.UUID Now time.Time AgentConfigSnapshot *executionstore.AgentConfigSnapshotRecord + ContextData *executionstore.ModelContextData CheckpointOverride *CheckpointRef MediaProjector MediaProjector } diff --git a/internal/modelprotocol/retry.go b/internal/modelprotocol/retry.go new file mode 100644 index 000000000..dae74e8cd --- /dev/null +++ b/internal/modelprotocol/retry.go @@ -0,0 +1,30 @@ +package modelprotocol + +import ( + "fmt" + "hash/fnv" + "time" +) + +func RetryBackoff(attemptNumber int, contextID string) time.Duration { + if attemptNumber < 1 { + attemptNumber = 1 + } + delay := time.Second + for i := 1; i < attemptNumber && delay < 30*time.Second; i++ { + delay *= 2 + } + + percent := deterministicModelCallRetryJitterPercent(contextID, attemptNumber) + delay = time.Duration(int64(delay) * int64(percent) / 100) + if delay > 30*time.Second { + return 30 * time.Second + } + return delay +} + +func deterministicModelCallRetryJitterPercent(contextID string, attemptNumber int) int { + hash := fnv.New32a() + _, _ = hash.Write([]byte(fmt.Sprintf("%s:%d", contextID, attemptNumber))) + return 80 + int(hash.Sum32()%41) +} diff --git a/internal/modelretry/policy.go b/internal/modelretry/policy.go index 4ea9bee7a..bd1c563e2 100644 --- a/internal/modelretry/policy.go +++ b/internal/modelretry/policy.go @@ -9,6 +9,7 @@ import ( "unicode/utf8" "github.com/omnara-ai/omnara/internal/model" + "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/storage/executionstore" ) @@ -258,7 +259,7 @@ func deterministicInputFailure(kind model.ErrorKind) bool { } func backoff(attemptNumber int, contextID string) time.Duration { - return executionstore.ModelCallRetryBackoff(attemptNumber, contextID) + return modelprotocol.RetryBackoff(attemptNumber, contextID) } func errorText(err error) string { diff --git a/internal/notifications/notifications.go b/internal/notifications/notifications.go index 3034126ef..1fc8a14d8 100644 --- a/internal/notifications/notifications.go +++ b/internal/notifications/notifications.go @@ -5,6 +5,8 @@ import ( "encoding/json" "errors" "fmt" + "maps" + "slices" "github.com/google/uuid" ) @@ -382,3 +384,36 @@ func (n *TxNotifications) Flush(ctx context.Context, publisher PostCommitPublish publisher.PublishPostCommit(ctx, intent) } } + +func (n *TxNotifications) Clone() *TxNotifications { + cloned := *n + cloned.daemonWorkByMachine = maps.Clone(n.daemonWorkByMachine) + cloned.runtimeEndedByID = maps.Clone(n.runtimeEndedByID) + cloned.processTerminationByMachine = make(map[uuid.UUID]map[uuid.UUID]struct{}, len(n.processTerminationByMachine)) + for id, processes := range n.processTerminationByMachine { + cloned.processTerminationByMachine[id] = maps.Clone(processes) + } + cloned.agentEventByID = make(map[uuid.UUID][]AgentEventReference, len(n.agentEventByID)) + for id, events := range n.agentEventByID { + cloned.agentEventByID[id] = slices.Clone(events) + } + cloned.toolCallUpdates = slices.Clone(n.toolCallUpdates) + for i, update := range cloned.toolCallUpdates { + if update.InteractionUpdate != nil { + interaction := *update.InteractionUpdate + cloned.toolCallUpdates[i].InteractionUpdate = &interaction + } + } + cloned.workerControls = slices.Clone(n.workerControls) + for i, control := range cloned.workerControls { + if control.Control.Cancel != nil { + cancel := *control.Control.Cancel + cloned.workerControls[i].Control.Cancel = &cancel + } + } + return &cloned +} + +func (n *TxNotifications) Restore(snapshot *TxNotifications) { + *n = *snapshot.Clone() +} diff --git a/internal/notifications/notifications_snapshot_test.go b/internal/notifications/notifications_snapshot_test.go new file mode 100644 index 000000000..6cb0fa7ae --- /dev/null +++ b/internal/notifications/notifications_snapshot_test.go @@ -0,0 +1,34 @@ +package notifications + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestTxNotificationsSnapshotOwnsNestedValues(t *testing.T) { + n := NewTxNotifications() + machine, process, runtime, agent, worker := uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New() + n.AddDaemonWork(machine) + n.AddDaemonProcessTermination(machine, process) + n.AddDaemonRuntimeEnded(runtime, machine, DaemonRuntimeEndReconnect) + n.AddAgentEvent(agent, 1, "agent_input") + n.AddToolCallUpdate(agent, uuid.New(), "waiting", &AgentInteractionUpdate{ID: uuid.New(), State: "open"}) + n.AddWorkerControlCancel(worker, agent, runtime) + snapshot := n.Clone() + expected := snapshot.Clone() + delete(n.daemonWorkByMachine, machine) + delete(n.processTerminationByMachine[machine], process) + delete(n.runtimeEndedByID, runtime) + n.agentEventByID[agent][0].Sequence = 2 + n.toolCallUpdates[0].InteractionUpdate.State = "resolved" + n.workerControls[0].Control.Cancel.RuntimeLockID = uuid.New() + require.Equal(t, expected, snapshot) + n.Restore(snapshot) + require.Equal(t, expected, n) + snapshot.toolCallUpdates[0].InteractionUpdate.State = "canceled" + snapshot.agentEventByID[agent][0].Sequence = 3 + delete(snapshot.processTerminationByMachine[machine], process) + require.Equal(t, expected, n) +} diff --git a/internal/processresult/results.go b/internal/processresult/results.go new file mode 100644 index 000000000..69ba107c7 --- /dev/null +++ b/internal/processresult/results.go @@ -0,0 +1,299 @@ +package processresult + +import ( + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/processaction" + "github.com/omnara-ai/omnara/internal/processcmd" + "github.com/omnara-ai/omnara/internal/publicid" +) + +type Process struct { + ID uuid.UUID + State string + DefaultOutputCursor int64 + SourceStartedAt *time.Time + SourceEndedAt *time.Time + ExitCode *int + ExitSignal string + StateReasonCode string + StateReasonMessage string + ExecutionSpec processcmd.ExecutionSpec +} +type Action struct { + ID uuid.UUID + Payload json.RawMessage +} +type daemonProcessReadObservation struct { + ProcessID string `json:"process_id"` + Output *string `json:"output"` + Cursor *int64 `json:"cursor"` + NextCursor *int64 `json:"next_cursor"` + Truncated *bool `json:"truncated"` +} + +func TerminalState(state string) bool { + switch state { + case "exited", "failed", "killed", "unknown": + return true + default: + return false + } +} +func publicID(kind publicid.Kind, id uuid.UUID) string { s, _ := publicid.Encode(kind, id); return s } +func Read( + process Process, + action Action, + raw json.RawMessage, +) (json.RawMessage, *int64, error) { + payload, err := processaction.DecodeReadPayload(action.Payload) + if err != nil { + return nil, nil, fmt.Errorf("decode read request: %w", err) + } + + var observed daemonProcessReadObservation + if err := json.Unmarshal(raw, &observed); err != nil { + return nil, nil, fmt.Errorf("decode read observation: %w", err) + } + if observed.Output == nil || + observed.Cursor == nil || + observed.NextCursor == nil || + observed.Truncated == nil { + return nil, nil, errors.New("read observation is incomplete") + } + publicProcessID := publicID(publicid.KindProcess, process.ID) + if observed.ProcessID != publicProcessID { + return nil, nil, errors.New("read observation names a different process") + } + minimumCursor := process.DefaultOutputCursor + if payload.Cursor != nil { + minimumCursor = *payload.Cursor + } + if *observed.Cursor < 0 || + *observed.Cursor < minimumCursor || + (*observed.Cursor > minimumCursor && !*observed.Truncated) || + *observed.NextCursor < *observed.Cursor { + return nil, nil, errors.New("read observation cursor range is invalid") + } + span := *observed.NextCursor - *observed.Cursor + maxBytes := payload.ObservationBytes() + if span > int64(maxBytes) || + len([]byte(*observed.Output)) > maxBytes || + int64(len([]byte(*observed.Output))) > span || + (span > 0 && *observed.Output == "") { + return nil, nil, errors.New("read observation exceeds its requested range") + } + + result := map[string]any{ + "process_id": publicProcessID, + "state": process.State, + "output": *observed.Output, + "cursor": *observed.Cursor, + "next_cursor": *observed.NextCursor, + "truncated": *observed.Truncated, + "done": TerminalState(process.State), + } + if process.SourceStartedAt != nil { + result["started_at"] = *process.SourceStartedAt + } + if TerminalState(process.State) { + result["ended_at"] = process.SourceEndedAt + result["exit_code"] = process.ExitCode + result["exit_signal"] = process.ExitSignal + result["state_reason_code"] = process.StateReasonCode + result["state_reason_message"] = process.StateReasonMessage + } + canonical, err := json.Marshal(result) + if err != nil { + return nil, nil, err + } + if payload.Cursor != nil { + return canonical, nil, nil + } + next := *observed.NextCursor + return canonical, &next, nil +} + +func ReadFailure( + process Process, + action Action, + code string, + message string, +) (json.RawMessage, error) { + result := map[string]any{ + "process_id": publicID(publicid.KindProcess, process.ID), + "process_action_id": publicID(publicid.KindProcessAction, action.ID), + "state": process.State, + "error_code": code, + "message": message, + "error": message, + "retryable": code == "machine_unreachable", + "done": TerminalState(process.State), + } + if process.SourceStartedAt != nil { + result["started_at"] = *process.SourceStartedAt + } + if TerminalState(process.State) { + result["ended_at"] = process.SourceEndedAt + result["exit_code"] = process.ExitCode + result["exit_signal"] = process.ExitSignal + result["state_reason_code"] = process.StateReasonCode + result["state_reason_message"] = process.StateReasonMessage + } + return json.Marshal(result) +} + +func Terminal(process Process) (string, json.RawMessage, error) { + if !TerminalState(process.State) { + return "", nil, fmt.Errorf("process %s is not terminal", process.ID) + } + exitCode := any(nil) + if process.ExitCode != nil { + exitCode = *process.ExitCode + } + errText := process.StateReasonMessage + if errText == "" { + errText = process.StateReasonCode + } + succeeded := process.State == "exited" && process.ExitCode != nil && *process.ExitCode == 0 + if process.State == "exited" && !succeeded && errText == "" && process.ExitCode != nil { + errText = fmt.Sprintf("exit code %d", *process.ExitCode) + } + result, err := json.Marshal(map[string]any{ + "process_id": publicID(publicid.KindProcess, process.ID), + "exit_code": exitCode, + "state": process.State, + "state_reason_code": process.StateReasonCode, + "error": errText, + }) + if err != nil { + return "", nil, fmt.Errorf("marshal process tool result: %w", err) + } + outcome := "succeeded" + if !succeeded { + outcome = "failed" + } + return outcome, result, nil +} + +func Started(process Process, observed json.RawMessage) (json.RawMessage, error) { + commandLabel := process.ExecutionSpec.Label() + result := map[string]any{ + "process_id": publicID(publicid.KindProcess, process.ID), + "state": process.State, + "command": commandLabel, + "next_action": "use write_process for stdin or read_process for retained output", + } + if len(observed) != 0 && string(observed) != "null" && string(observed) != "{}" { + var observedObject map[string]any + if err := json.Unmarshal(observed, &observedObject); err != nil { + return nil, fmt.Errorf("decode process started result: %w", err) + } + for key, value := range observedObject { + result[key] = value + } + result["process_id"] = publicID(publicid.KindProcess, process.ID) + result["state"] = process.State + result["command"] = commandLabel + result["next_action"] = "use write_process for stdin or read_process for retained output" + } + return json.Marshal(result) +} + +func ActionResult( + processID uuid.UUID, + actionID uuid.UUID, + state string, + reasonCode string, + errText string, +) (json.RawMessage, error) { + body, err := json.Marshal(struct { + ProcessID string `json:"process_id"` + ProcessActionID string `json:"process_action_id"` + State string `json:"state"` + StateReasonCode string `json:"state_reason_code"` + Error string `json:"error"` + }{ + ProcessID: publicID(publicid.KindProcess, processID), + ProcessActionID: publicID(publicid.KindProcessAction, actionID), + State: state, + StateReasonCode: reasonCode, + Error: errText, + }) + if err != nil { + return nil, fmt.Errorf("marshal process action tool result: %w", err) + } + return body, nil +} + +func ReadMatches( + raw json.RawMessage, + result json.RawMessage, +) (bool, error) { + var observed daemonProcessReadObservation + ok := json.Unmarshal(raw, &observed) == nil + if !ok { + return false, nil + } + var committed daemonProcessReadObservation + if err := json.Unmarshal(result, &committed); err != nil { + return false, fmt.Errorf("decode committed read result: %w", err) + } + if observed.Output == nil || committed.Output == nil || + observed.Cursor == nil || committed.Cursor == nil || + observed.NextCursor == nil || committed.NextCursor == nil || + observed.Truncated == nil || committed.Truncated == nil { + return false, nil + } + return observed.ProcessID == committed.ProcessID && + *observed.Output == *committed.Output && + *observed.Cursor == *committed.Cursor && + *observed.NextCursor == *committed.NextCursor && + *observed.Truncated == *committed.Truncated, nil +} + +func ObservationNextCursor( + raw json.RawMessage, +) (int64, bool, error) { + if len(raw) == 0 || string(raw) == "null" || string(raw) == "{}" { + return 0, false, nil + } + var observed struct { + Output *string `json:"output"` + Cursor *int64 `json:"cursor"` + NextCursor *int64 `json:"next_cursor"` + } + if err := json.Unmarshal(raw, &observed); err != nil { + return 0, false, err + } + if observed.NextCursor == nil { + if observed.Cursor != nil || observed.Output != nil { + return 0, false, errors.New( + "process observation cursor range is incomplete", + ) + } + return 0, false, nil + } + if observed.Cursor == nil || observed.Output == nil || + *observed.Cursor < 0 || + *observed.NextCursor < *observed.Cursor { + return 0, false, errors.New( + "process observation cursor range is invalid", + ) + } + span := *observed.NextCursor - *observed.Cursor + outputBytes := len([]byte(*observed.Output)) + if span > int64(processaction.MaxObservationBytes) || + outputBytes > processaction.MaxObservationBytes || + int64(outputBytes) > span || + (span > 0 && *observed.Output == "") { + return 0, false, errors.New( + "process observation exceeds its allowed range", + ) + } + return *observed.NextCursor, true, nil +} diff --git a/internal/storage/executionstore/process_read_results_test.go b/internal/processresult/results_test.go similarity index 88% rename from internal/storage/executionstore/process_read_results_test.go rename to internal/processresult/results_test.go index a079f1d74..e6bb4e903 100644 --- a/internal/storage/executionstore/process_read_results_test.go +++ b/internal/processresult/results_test.go @@ -1,4 +1,4 @@ -package executionstore +package processresult import ( "encoding/json" @@ -11,12 +11,12 @@ import ( func TestCanonicalProcessReadResultValidatesSkippedOutput(t *testing.T) { t.Parallel() - process := ProcessRecord{ + process := Process{ ID: uuid.New(), - State: ProcessStateRunning, + State: "running", DefaultOutputCursor: 4, } - publicProcessID := publicResourceID(publicid.KindProcess, process.ID) + publicProcessID := publicID(publicid.KindProcess, process.ID) tests := []struct { name string payload json.RawMessage @@ -56,9 +56,9 @@ func TestCanonicalProcessReadResultValidatesSkippedOutput(t *testing.T) { t.Run(test.name, func(t *testing.T) { t.Parallel() - _, advance, err := canonicalProcessReadResult( + _, advance, err := Read( process, - ProcessActionRecord{Payload: test.payload}, + Action{Payload: test.payload}, json.RawMessage(test.observation), ) if (err != nil) != test.wantError { diff --git a/internal/storage/artifactstore/artifacts_store.go b/internal/storage/artifactstore/artifacts_store.go index eaa4a4189..39bf86d50 100644 --- a/internal/storage/artifactstore/artifacts_store.go +++ b/internal/storage/artifactstore/artifacts_store.go @@ -293,7 +293,7 @@ func (s *Store) ListAgentArtifactsByIDs( func insertArtifactTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, artifactID uuid.UUID, input CreateArtifactInput, ) (ArtifactRecord, error) { diff --git a/internal/storage/artifactstore/prepared_artifacts.go b/internal/storage/artifactstore/prepared_artifacts.go index e1d27170c..de5901013 100644 --- a/internal/storage/artifactstore/prepared_artifacts.go +++ b/internal/storage/artifactstore/prepared_artifacts.go @@ -44,7 +44,7 @@ func (p PreparedArtifact) Validate() error { func InsertPreparedArtifactsTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID uuid.UUID, artifacts []PreparedArtifact, ) ([]ArtifactRecord, error) { diff --git a/internal/storage/default_reconciliation_integration_test.go b/internal/storage/default_reconciliation_integration_test.go index 4588b3951..4412a5663 100644 --- a/internal/storage/default_reconciliation_integration_test.go +++ b/internal/storage/default_reconciliation_integration_test.go @@ -14,10 +14,10 @@ import ( "github.com/jackc/pgx/v5/pgxpool" "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/modelprotocol" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/secrets" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/management" "github.com/omnara-ai/omnara/internal/storage/modelstore" @@ -166,13 +166,15 @@ func TestReconcileDefaults(t *testing.T) { Material: secrets.GenericMaterial{Value: "tenant-test-key"}, Actor: userPrincipal(user.ID), }) require.NoError(t, err) - tenantProvider, err := store.Models().CreateModelProviderConfig(ctx, modelstore.CreateModelProviderConfigInput{ - OrgID: created.Org.ID, Name: "tenant-timeouts", APIFormat: modelprotocol.APIFormatOpenAIResponses, - BaseURL: "https://example.test", CredentialSecretID: tenantCredential.ID, - RequestTimeoutMS: 600000, IdleTimeoutMS: 700000, - }) + tenantProvider, err := store.Models(). + CreateModelProviderConfig(ctx, modelstore.CreateModelProviderConfigInput{ + OrgID: created.Org.ID, Name: "tenant-timeouts", APIFormat: modelprotocol.APIFormatOpenAIResponses, + BaseURL: "https://example.test", CredentialSecretID: tenantCredential.ID, + RequestTimeoutMS: 600000, IdleTimeoutMS: 700000, + }) require.NoError(t, err) - updateModel, err := store.Models().GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "update-model") + updateModel, err := store.Models(). + GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "update-model") if err != nil { t.Fatalf("get update model: %v", err) } @@ -220,18 +222,21 @@ model: } return config } - removeModel, err := store.Models().GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "remove-model") + removeModel, err := store.Models(). + GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "remove-model") if err != nil { t.Fatalf("get remove model: %v", err) } historicalConfig := createAgentConfig(removeModel) - retainedModel, err := store.Models().GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "retained-model") + retainedModel, err := store.Models(). + GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "retained-model") if err != nil { t.Fatalf("get retained model: %v", err) } - manualProject, err := store.Identity().CreateProjectForPrincipal(ctx, identitystore.CreateProjectForPrincipalInput{ - OrgID: created.Org.ID, Creator: userPrincipal(user.ID), Name: "Manual", IdempotencyKey: "manual", - }) + manualProject, err := store.Identity(). + CreateProjectForPrincipal(ctx, identitystore.CreateProjectForPrincipalInput{ + OrgID: created.Org.ID, Creator: userPrincipal(user.ID), Name: "Manual", IdempotencyKey: "manual", + }) if err != nil { t.Fatalf("create manual project: %v", err) } @@ -360,7 +365,12 @@ model: ContextWindowTokens: 16384, MaxOutputTokens: new(2048), }, - {Name: "add-model", ProviderModelSlug: "example/add", ContextWindowTokens: 8192, MaxOutputTokens: new(1024)}, + { + Name: "add-model", + ProviderModelSlug: "example/add", + ContextWindowTokens: 8192, + MaxOutputTokens: new(1024), + }, { Name: "tenant-model", ProviderModelSlug: "example/cluster-collision", @@ -375,7 +385,11 @@ model: assertRetainedModelWarnings := func(label string, result orglifecycle.ReconcileDefaultsResult) { t.Helper() if len(result.Warnings) != 3 { - t.Fatalf("%s warnings = %v, want tenant collision and two retained-model warnings", label, result.Warnings) + t.Fatalf( + "%s warnings = %v, want tenant collision and two retained-model warnings", + label, + result.Warnings, + ) } warnings := strings.Join(result.Warnings, "\n") for _, modelName := range []string{activeAgentModel.Name, profileModel.Name} { @@ -441,17 +455,21 @@ model: t.Fatal("runtime mismatch marker survived runtime protection change") } - reconciledProvider, err := store.Models().GetModelProviderConfigByName(ctx, created.Org.ID, desiredProvider.Name) + reconciledProvider, err := store.Models(). + GetModelProviderConfigByName(ctx, created.Org.ID, desiredProvider.Name) if err != nil { t.Fatalf("get reconciled provider: %v", err) } - if reconciledProvider.BaseURL != desiredProvider.BaseURL || reconciledProvider.RequestTimeoutMS != 120000 || + if reconciledProvider.BaseURL != desiredProvider.BaseURL || + reconciledProvider.RequestTimeoutMS != 120000 || reconciledProvider.IdleTimeoutMS != 45000 || reconciledProvider.CredentialSecretID != provider.CredentialSecretID { t.Fatalf("unexpected reconciled provider: %+v", reconciledProvider) } - updatedModel, err := store.Models().GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "update-model") - if err != nil || updatedModel.DefaultMaxOutputTokens != nil || updatedModel.ProviderModelSlug != "example/new" || + updatedModel, err := store.Models(). + GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "update-model") + if err != nil || updatedModel.DefaultMaxOutputTokens != nil || + updatedModel.ProviderModelSlug != "example/new" || (updatedModel.MaxOutputTokens == nil || *updatedModel.MaxOutputTokens != 2048) { t.Fatalf("unexpected updated model: %+v, err %v", updatedModel, err) } @@ -607,7 +625,8 @@ model: t.Fatalf("machine pool description = %q, want %q", poolRecord.Description, desiredPool.Description) } assertDecodedJSONEqual(t, poolRecord.DefaultMachineSecretEnv, string(organizationSecretEnv)) - updatedModel, err = store.Models().GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "update-model") + updatedModel, err = store.Models(). + GetConfiguredModelByName(ctx, created.Org.ID, provider.ID, "update-model") if err != nil || updatedModel.ProviderModelSlug != "example/without-project" || updatedModel.MaxOutputTokens != nil { t.Fatalf("unexpected model updated without default project: %+v, err %v", updatedModel, err) @@ -992,20 +1011,19 @@ INSERT INTO machines( }) integrationdb.WaitForNamedLockWaiters(t, ctx, pool, "LockMachineForLifecycle", 1) deleteDone := integrationdb.RunAsyncError(func() error { - deleteTx, deleteErr := pool.Begin(ctx) + deleteTx, deleteErr := agentexecution.NewCell("primary", pool, nil, nil).Begin(ctx) if deleteErr != nil { return deleteErr } defer func() { _ = deleteTx.Rollback(ctx) }() - _, deleteErr = store.Execution().DeleteMachinePoolTx( + _, deleteErr = store.Execution().DeleteMachinePoolInUnit( ctx, deleteTx, - notifications.NewTxNotifications(), created.Org.ID, poolBRow.ID, ) if deleteErr == nil { - deleteErr = deleteTx.Commit(ctx) + deleteErr = deleteTx.Commit(ctx, "delete pool") } return deleteErr }) @@ -1184,13 +1202,14 @@ machine_sources: }) integrationdb.WaitForNamedLockWaiters(t, ctx, pool, "LockMachinePoolForLifecycle", 1) changeDone := integrationdb.RunAsyncError(func() error { - _, changeErr := store.Execution().IntegrationChangeAgentConfigOnce(ctx, executionstore.ChangeAgentConfigInput{ - CreateAgentConfigInput: nextConfigInput, - AgentID: launched.Agent.ID, - ActorType: identitystore.PrincipalTypeUser, - ActorID: user.ID, - IdempotencyKey: "reconcile-config-order", - }) + _, changeErr := store.Execution(). + IntegrationChangeAgentConfigOnce(ctx, executionstore.ChangeAgentConfigInput{ + CreateAgentConfigInput: nextConfigInput, + AgentID: launched.Agent.ID, + ActorType: identitystore.PrincipalTypeUser, + ActorID: user.ID, + IdempotencyKey: "reconcile-config-order", + }) return changeErr }) waitForDefaultReconciliationLockWaiter( diff --git a/internal/storage/executionstore/actors_store.go b/internal/storage/executionstore/actors_store.go index c2a3cee76..834228815 100644 --- a/internal/storage/executionstore/actors_store.go +++ b/internal/storage/executionstore/actors_store.go @@ -164,7 +164,7 @@ func (s *Store) PutActor(ctx context.Context, input PutActorInput) (ActorRecord, func enterActiveActorProjectTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, qtx *dbsqlc.Queries, projectID uuid.UUID, ) error { diff --git a/internal/storage/executionstore/agent_config_change_gates_integration_test.go b/internal/storage/executionstore/agent_config_change_gates_integration_test.go index dec6015db..7d69c3c44 100644 --- a/internal/storage/executionstore/agent_config_change_gates_integration_test.go +++ b/internal/storage/executionstore/agent_config_change_gates_integration_test.go @@ -12,6 +12,7 @@ import ( "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -19,7 +20,9 @@ import ( "github.com/stretchr/testify/require" ) -func (f integrationActivationFixture) launchWithSelectedIntegration(t *testing.T) executionstore.LaunchAgentResult { +func (f integrationActivationFixture) launchWithSelectedIntegration( + t *testing.T, +) executionstore.LaunchAgentResult { t.Helper() definition := f.withSendingTools(t, f.definition(t, "Selected integration")) var compiled agentconfig.Compiled @@ -38,24 +41,27 @@ func (f integrationActivationFixture) launchWithSelectedIntegration(t *testing.T launch.Agent.ID, "C123:111.222", ) - tx := integrationdb.BeginTx(t, f.ctx, f.store.pool) + txUnit, unitErr := f.store.Execution().IntegrationBeginUnit(f.ctx) + require.NoError(t, unitErr) + defer func() { _ = txUnit.Rollback(f.ctx) }() + tx := txUnit.DB() require.NoError(t, lifecyclelock.EnterActiveProject(f.ctx, tx, testOrgID, testProjectID)) require.NoError(t, integrationstore.LockIntegrationsTx(f.ctx, tx, testProjectID, nil, f.integration.ID)) - _, err = dbsqlc.New(tx).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ + _, err = txUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: launch.Agent.ID, - }) + }, agentexecution.ExternalAuthority{}) require.NoError(t, err) require.NoError(t, f.store.Integrations().AssignAgentIntegrationConversationTx( f.ctx, tx, testProjectID, launch.Agent.ID, f.integration.ID, integrationstore.ConversationAddress{Kind: target.ScopeKind, Ref: target.ScopeRef}, )) selection, err := executionstore.IntegrationSelectInteractionDestinationForOriginTx( - f.ctx, tx, testProjectID, launch.Agent.ID, target.ID, + f.ctx, txUnit, testProjectID, launch.Agent.ID, target.ID, ) require.NoError(t, err) require.Equal(t, "chat", selection.HandlerKey) require.Equal(t, target.ID, selection.IntegrationTargetID) - require.NoError(t, tx.Commit(f.ctx)) + require.NoError(t, txUnit.Commit(f.ctx, "fixture admission")) return launch } @@ -124,9 +130,21 @@ func TestConfigChangeSerializesNextIntegrationWithRevocation(t *testing.T) { f.ctx, dbsqlc.LockAgentMachineSourcesParams{AgentID: launch.Agent.ID}, )) changed = integrationdb.RunAsync(change) - integrationdb.WaitForNamedLockWaiters(t, f.ctx, f.store.pool, "LockAgentMachineSources", 1) + integrationdb.WaitForNamedLockWaiters( + t, + f.ctx, + f.store.pool, + "LockAgentMachineSources", + 1, + ) revoked = integrationdb.RunAsyncError(revoke) - integrationdb.WaitForNamedLockWaiters(t, f.ctx, f.store.pool, "LockIntegrationLifecycleExclusive", 1) + integrationdb.WaitForNamedLockWaiters( + t, + f.ctx, + f.store.pool, + "LockIntegrationLifecycleExclusive", + 1, + ) } else { require.NoError(t, dbsqlc.New(control).LockIntegrationLifecycleShared( f.ctx, dbsqlc.LockIntegrationLifecycleSharedParams{IntegrationID: f.integration.ID}, @@ -137,19 +155,36 @@ func TestConfigChangeSerializesNextIntegrationWithRevocation(t *testing.T) { integrationdb.WaitForNamedLockWaiters(t, f.ctx, f.store.pool, "LockIntegrationLifecycleShared", 1) } require.NoError(t, control.Commit(f.ctx)) - result := integrationdb.AwaitSuccess(t, changed, "config change racing integration revocation") - require.NoError(t, integrationdb.Await(t, revoked, "integration revocation racing config change")) + result := integrationdb.AwaitSuccess( + t, + changed, + "config change racing integration revocation", + ) + require.NoError( + t, + integrationdb.Await(t, revoked, "integration revocation racing config change"), + ) current, err := f.store.Execution().GetAgentInProject(f.ctx, testProjectID, launch.Agent.ID) require.NoError(t, err) - require.Equal(t, result.AgentConfig.ID, current.CurrentConfigID, "revoked references remain valid metadata") + require.Equal( + t, + result.AgentConfig.ID, + current.CurrentConfigID, + "revoked references remain valid metadata", + ) captureTx := integrationdb.BeginTx(t, f.ctx, f.store.pool) destination, err := executionstore.IntegrationCaptureInteractionDestinationTx( f.ctx, captureTx, testProjectID, launch.Agent.ID, ) require.NoError(t, err) - require.Nil(t, destination, "config activation cannot grant live authority to a revoked integration") + require.Nil( + t, + destination, + "config activation cannot grant live authority to a revoked integration", + ) if state == "deleted" || order == "revocation-first" { - selection, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, launch.Agent.ID) + selection, err := f.store.Execution(). + GetInteractionSelection(f.ctx, testProjectID, launch.Agent.ID) require.NoError(t, err) require.Equal(t, executionstore.InteractionSelection{AutoSelect: true}, selection) } @@ -170,7 +205,12 @@ func TestConfigChangeReplaySkipsRevokedIntegrationGate(t *testing.T) { t.Parallel() f := newIntegrationActivationFixture(t) launch := f.launchWithSelectedIntegration(t) - input := f.changeInput(t, launch.Agent.ID, "Accepted integration config", "accepted-integration-config") + input := f.changeInput( + t, + launch.Agent.ID, + "Accepted integration config", + "accepted-integration-config", + ) input.CreateAgentConfigInput = f.withSendingTools(t, input.CreateAgentConfigInput) input.ExpectedCurrentConfigID = launch.Agent.CurrentConfigID accepted, err := f.store.Execution().ChangeAgentConfig(f.ctx, input) @@ -188,7 +228,11 @@ func TestConfigChangeReplaySkipsRevokedIntegrationGate(t *testing.T) { ctx, cancel := context.WithTimeout(f.ctx, 2*time.Second) defer cancel() replayed, err := f.store.Execution().IntegrationChangeAgentConfigOnce(ctx, input) - require.NoError(t, err, "committed replay skips integration gates and stale expected-current validation") + require.NoError( + t, + err, + "committed replay skips integration gates and stale expected-current validation", + ) require.Equal(t, accepted.AgentConfig.ID, replayed.AgentConfig.ID) require.Equal(t, accepted.ConfigChange.AgentInput.ID, replayed.ConfigChange.AgentInput.ID) require.Equal(t, accepted.ConfigChange.Event.ID, replayed.ConfigChange.Event.ID) @@ -197,9 +241,20 @@ func TestConfigChangeReplaySkipsRevokedIntegrationGate(t *testing.T) { require.ErrorIs(t, err, storeerr.ErrIdempotencyConflict) current, err := f.store.Execution().GetAgentInProject(f.ctx, testProjectID, launch.Agent.ID) require.NoError(t, err) - require.Equal(t, latest.AgentConfig.ID, current.CurrentConfigID, "replay cannot reactivate the older config") - require.Equal(t, before, f.subscriptions(t, launch.Agent.ID), "replay cannot restore subscriptions") - selection, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, launch.Agent.ID) + require.Equal( + t, + latest.AgentConfig.ID, + current.CurrentConfigID, + "replay cannot reactivate the older config", + ) + require.Equal( + t, + before, + f.subscriptions(t, launch.Agent.ID), + "replay cannot restore subscriptions", + ) + selection, err := f.store.Execution(). + GetInteractionSelection(f.ctx, testProjectID, launch.Agent.ID) require.NoError(t, err) require.Equal(t, executionstore.InteractionSelection{AutoSelect: true}, selection) var events int diff --git a/internal/storage/executionstore/agent_config_changes_store.go b/internal/storage/executionstore/agent_config_changes_store.go index 8ecfb6d15..3850a1c85 100644 --- a/internal/storage/executionstore/agent_config_changes_store.go +++ b/internal/storage/executionstore/agent_config_changes_store.go @@ -7,11 +7,10 @@ import ( "reflect" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/events" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -49,7 +48,10 @@ type ChangeAgentConfigResult struct { DeleteMachines []MachineRecord } -func (s *Store) ChangeAgentConfig(ctx context.Context, input ChangeAgentConfigInput) (ChangeAgentConfigResult, error) { +func (s *Store) ChangeAgentConfig( + ctx context.Context, + input ChangeAgentConfigInput, +) (ChangeAgentConfigResult, error) { if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil { return ChangeAgentConfigResult{}, errors.New("project and agent are required") } @@ -62,13 +64,14 @@ func (s *Store) changeAgentConfigOnce( ctx context.Context, input ChangeAgentConfigInput, ) (ChangeAgentConfigResult, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return ChangeAgentConfigResult{}, fmt.Errorf("begin change agent config: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) project, err := loadProjectTx(ctx, qtx, input.ProjectID) if err != nil { return ChangeAgentConfigResult{}, err @@ -93,7 +96,9 @@ func (s *Store) changeAgentConfigOnce( return ChangeAgentConfigResult{}, fmt.Errorf("load agent for config change: %w", err) } if agent.ParentAgentID != nil { - return ChangeAgentConfigResult{}, storeerr.InvalidRequest(errors.New("subagent configurations are read-only")) + return ChangeAgentConfigResult{}, storeerr.InvalidRequest( + errors.New("subagent configurations are read-only"), + ) } idempotentReplay, err := configChangeReplayExistsTx(ctx, qtx, input) if err != nil { @@ -160,7 +165,10 @@ func (s *Store) changeAgentConfigOnce( }, ) if err != nil { - return ChangeAgentConfigResult{}, fmt.Errorf("list agent pool machines for config change: %w", err) + return ChangeAgentConfigResult{}, fmt.Errorf( + "list agent pool machines for config change: %w", + err, + ) } if err := lockLaunchMachineSourcesTx( ctx, @@ -172,7 +180,7 @@ func (s *Store) changeAgentConfigOnce( return ChangeAgentConfigResult{}, err } } - if err := lockAgentForConfigActivationTx(ctx, qtx, activationInput); err != nil { + if err := lockAgentForConfigActivationTx(ctx, unit, activationInput); err != nil { return ChangeAgentConfigResult{}, err } if err := authorizeAgentConfigChangeTx(ctx, qtx, activationInput); err != nil { @@ -190,8 +198,7 @@ func (s *Store) changeAgentConfigOnce( } configChange, err := activateLockedAuthorizedAgentConfigTx( ctx, - txNotifications, - tx, + unit, qtx, activationInput, ) @@ -208,13 +215,13 @@ func (s *Store) changeAgentConfigOnce( return ChangeAgentConfigResult{}, fmt.Errorf("reload agent after config change: %w", err) } if currentAgent.CurrentConfigID == config.ID { - if _, err := s.ReconcileInteractionSelectionTx(ctx, tx, input.ProjectID, input.AgentID); err != nil { + if _, err := s.ReconcileInteractionSelectionInUnit(ctx, unit, input.ProjectID, input.AgentID); err != nil { return ChangeAgentConfigResult{}, err } deleteMachines, err = s.reconcileAgentMachineSourcesTx( ctx, txNotifications, - tx, + unit, qtx, input.ProjectID, input.AgentID, @@ -240,14 +247,7 @@ func (s *Store) changeAgentConfigOnce( } } } - if err := qtx.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: []byte(`{"reason":"config_change"}`), - }); err != nil { - return ChangeAgentConfigResult{}, fmt.Errorf("reconcile config-change wakeup: %w", err) - } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "change agent config"); err != nil { + if err := unit.Commit(ctx, "change agent config"); err != nil { return ChangeAgentConfigResult{}, err } return ChangeAgentConfigResult{ @@ -296,29 +296,28 @@ func validateLiveAgentConfigChangeTx( func activateNewAgentConfigTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input ActivateAgentConfigInput, ) (AgentConfigChangeRecord, error) { if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.AgentConfigID == uuid.Nil { return AgentConfigChangeRecord{}, errors.New("project, agent, and config are required") } - return activateLockedAuthorizedAgentConfigTx(ctx, txNotifications, tx, qtx, input) + return activateLockedAuthorizedAgentConfigTx(ctx, unit, qtx, input) } func lockAgentForConfigActivationTx( ctx context.Context, - qtx *dbsqlc.Queries, + unit *agentexecution.Unit, input ActivateAgentConfigInput, ) error { if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.AgentConfigID == uuid.Nil { return errors.New("project, agent, and config are required") } - if _, err := qtx.LockAgentInProject( + if _, err := unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - ); err != nil { + agentexecution.IngressAuthority{}); err != nil { return fmt.Errorf("lock agent for config change: %w", err) } return nil @@ -347,8 +346,7 @@ func authorizeAgentConfigChangeTx( func activateLockedAuthorizedAgentConfigTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input ActivateAgentConfigInput, ) (AgentConfigChangeRecord, error) { @@ -373,90 +371,45 @@ func activateLockedAuthorizedAgentConfigTx( return AgentConfigChangeRecord{}, err } default: - return AgentConfigChangeRecord{}, fmt.Errorf("unsupported agent config change actor type %q", input.ActorType) + return AgentConfigChangeRecord{}, fmt.Errorf( + "unsupported agent config change actor type %q", + input.ActorType, + ) } - metadata, err := marshalJSON(map[string]any{"agent_config_id": input.AgentConfigID, "reason": input.Reason}) + metadata, err := marshalJSON( + map[string]any{"agent_config_id": input.AgentConfigID, "reason": input.Reason}, + ) if err != nil { return AgentConfigChangeRecord{}, fmt.Errorf("marshal config change metadata: %w", err) } - row, err := qtx.InsertConfigChangeAgentInput(ctx, dbsqlc.InsertConfigChangeAgentInputParams{ - AgentConfigID: storeutil.IDFromNil(input.AgentConfigID), - ActorID: storeutil.IDFromNil(actorID), - IdempotencyScope: storeutil.TextFromEmpty("agent_config_change"), - InputIdempotencyKey: storeutil.TextFromEmpty(input.IdempotencyKey), - Metadata: metadata, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }) + h, err := unit.Handle(input.ProjectID, input.AgentID) if err != nil { - return AgentConfigChangeRecord{}, fmt.Errorf("insert config-change agent input: %w", err) - } - agentInput := agentInputRecordFromConfigChangeSQLC(row) - if agentInput.ActorID != actorID || - agentInput.AgentConfigID != input.AgentConfigID { - return AgentConfigChangeRecord{}, storeerr.ErrIdempotencyConflict - } - if !sameJSON(agentInput.Metadata, metadata) { - return AgentConfigChangeRecord{}, storeerr.ErrIdempotencyConflict + return AgentConfigChangeRecord{}, err } - if agentInput.State == "resolved" { - eventRow, err := qtx.GetEventByProjectAgentIdempotencyKey( - ctx, - dbsqlc.GetEventByProjectAgentIdempotencyKeyParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - IdempotencyKey: "agent_input:" + agentInput.ID.String(), - }, - ) - if err != nil { - return AgentConfigChangeRecord{}, fmt.Errorf("load idempotent config-change event: %w", err) - } - event, err := eventFromProjectIdempotencySQLC(eventRow) - if err != nil { - return AgentConfigChangeRecord{}, err - } - if agentInput.AdmittedEventID != event.ID { - return AgentConfigChangeRecord{}, storeerr.ErrStateTransitionConflict - } - return AgentConfigChangeRecord{AgentInput: agentInput, Event: event}, nil + result, err := h.ActivateConfig(ctx, agentexecution.ActivateConfigInput{ + ConfigID: input.AgentConfigID, ActorID: actorID, IdempotencyKey: input.IdempotencyKey, Metadata: metadata}) + if err != nil { + return AgentConfigChangeRecord{}, err } - eventRecord, _, _, err := appendEventToCurrentOrNewAgentTurnTx( + row, err := qtx.GetAgentInput( ctx, - txNotifications, - tx, - qtx, - AppendTypedAgentEventInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Kind: events.KindAgentInput, - IdempotencyKey: "agent_input:" + agentInput.ID.String(), - AgentInputID: agentInput.ID, - }, - true, + dbsqlc.GetAgentInputParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: result.InputID}, ) if err != nil { return AgentConfigChangeRecord{}, err } - resolved, err := qtx.ResolveConfigChangeAgentInput( + eventRow, err := qtx.GetEventByProjectAgentIdempotencyKey( ctx, - dbsqlc.ResolveConfigChangeAgentInputParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: agentInput.ID, - EventID: eventRecord.Event.ID, - }, + dbsqlc.GetEventByProjectAgentIdempotencyKeyParams{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + IdempotencyKey: "agent_input:" + result.InputID.String()}, ) - if errors.Is(err, pgx.ErrNoRows) { - return AgentConfigChangeRecord{}, storeerr.ErrStateTransitionConflict - } if err != nil { - return AgentConfigChangeRecord{}, fmt.Errorf("resolve config-change agent input: %w", err) + return AgentConfigChangeRecord{}, err } - agentInput.State = "resolved" - agentInput.AdmittedEventID = eventRecord.Event.ID - agentInput.AdmittedAt = resolved.AdmittedAt - agentInput.ResolvedAt = resolved.ResolvedAt - return AgentConfigChangeRecord{AgentInput: agentInput, Event: eventRecord.Event}, nil + event, err := eventFromProjectIdempotencySQLC(eventRow) + return AgentConfigChangeRecord{AgentInput: agentInputRecordFromGetSQLC(row), Event: event}, err } func validateProjectPrincipalActionTx( @@ -489,28 +442,3 @@ func validateProjectPrincipalActionTx( } return storeerr.ErrUnauthorized } - -func agentInputRecordFromConfigChangeSQLC(row dbsqlc.InsertConfigChangeAgentInputRow) AgentInputRecord { - return AgentInputRecord{ - ID: row.ID, - ProjectID: row.ProjectID, - AgentID: row.AgentID, - State: row.State, - InputRank: row.InputRank, - ActorID: storeutil.IDFromPtr(row.ActorID), - InputKind: row.InputKind, - IdempotencyScope: row.IdempotencyScope, - InputIdempotencyKey: row.InputIdempotencyKey, - QueuedAt: row.QueuedAt, - AdmittedEventID: storeutil.IDFromPtr(row.AdmittedEventID), - AdmittedAt: row.AdmittedAt, - CanceledAt: row.CanceledAt, - DeliveryMode: AgentInputDeliveryMode(row.DeliveryMode), - ControlType: row.ControlType, - TargetInteractionID: storeutil.IDFromPtr(row.TargetInteractionID), - AgentConfigID: storeutil.IDFromPtr(row.AgentConfigID), - ResolvedAt: row.ResolvedAt, - RejectedReason: row.RejectedReason, - Metadata: row.Metadata, - } -} diff --git a/internal/storage/executionstore/agent_configs_store.go b/internal/storage/executionstore/agent_configs_store.go index 9f434e4cb..5e2925a41 100644 --- a/internal/storage/executionstore/agent_configs_store.go +++ b/internal/storage/executionstore/agent_configs_store.go @@ -115,6 +115,21 @@ func (s *Store) CaptureAgentConfigForModelContext( err, ) } + snapshot, err := captureAgentConfigForModelContextTx(ctx, qtx, projectID, agentID) + if err != nil { + return AgentConfigSnapshotRecord{}, err + } + if err := tx.Commit(ctx); err != nil { + return AgentConfigSnapshotRecord{}, fmt.Errorf("commit capture agent config for model context: %w", err) + } + return snapshot, nil +} + +func captureAgentConfigForModelContextTx( + ctx context.Context, + qtx *dbsqlc.Queries, + projectID, agentID uuid.UUID, +) (AgentConfigSnapshotRecord, error) { // The lock and the capture are deliberately separate statements. Under // READ COMMITTED, the second statement gets a fresh snapshot after any // transaction that previously held the agent lock has committed, so the @@ -133,12 +148,6 @@ func (s *Store) CaptureAgentConfigForModelContext( err, ) } - if err := tx.Commit(ctx); err != nil { - return AgentConfigSnapshotRecord{}, fmt.Errorf( - "commit capture agent config for model context: %w", - err, - ) - } return agentConfigSnapshotFromSQLC(row), nil } diff --git a/internal/storage/executionstore/agent_content_inputs_store.go b/internal/storage/executionstore/agent_content_inputs_store.go index 8b620a687..3849045a3 100644 --- a/internal/storage/executionstore/agent_content_inputs_store.go +++ b/internal/storage/executionstore/agent_content_inputs_store.go @@ -7,10 +7,9 @@ import ( "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/resourcemeta" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -58,20 +57,20 @@ func (s *Store) CreateAgentContentInput( if err != nil { return AgentInputRecord{}, nil, false, err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentInputRecord{}, nil, false, fmt.Errorf( "begin create agent content input: %w", err, ) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + qtx := dbsqlc.New(tx) result, err := createAgentContentInputTx( ctx, - txNotifications, - tx, + unit, qtx, input, contentBlocks, @@ -79,12 +78,7 @@ func (s *Store) CreateAgentContentInput( if err != nil { return AgentInputRecord{}, nil, false, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "create agent content input", - ); err != nil { + if err := unit.Commit(ctx, "create agent content input"); err != nil { return AgentInputRecord{}, nil, false, err } return result.agentInput, result.contentBlocks, result.created, nil @@ -106,7 +100,7 @@ func prepareCreateAgentContentInput( "cancel_open_interactions is allowed only for steering inputs", )) } - if input.IdempotencyScope == "" { + if input.IdempotencyKey != "" && input.IdempotencyScope == "" { input.IdempotencyScope = "content_input" } return input, nil @@ -121,15 +115,21 @@ type createAgentContentInputTxResult struct { func createAgentContentInputTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input CreateAgentContentInputInput, contentBlocks []CreateContentBlockInput, ) (createAgentContentInputTxResult, error) { - agent, err := qtx.LockAgentInProject( + var err error + input, err = prepareCreateAgentContentInput(input) + if err != nil { + return createAgentContentInputTxResult{}, err + } + tx := unit.DB() + agent, err := unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, + agentexecution.IngressAuthority{}, ) if err != nil { return createAgentContentInputTxResult{}, fmt.Errorf("lock agent for content input: %w", err) @@ -204,78 +204,35 @@ func createAgentContentInputTx( if err != nil { return createAgentContentInputTxResult{}, err } - agentInput, err := insertAgentInputTx(ctx, tx, insertAgentInputInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - DeliveryMode: input.DeliveryMode, - ActorID: actorID, - IntegrationTargetID: input.IntegrationTargetID, - IdempotencyScope: input.IdempotencyScope, - InputIdempotencyKey: input.IdempotencyKey, - Metadata: input.Metadata, - }) + h, err := unit.Handle(input.ProjectID, input.AgentID) if err != nil { return createAgentContentInputTxResult{}, err } - if err := createAgentInputContentBlocksTx( - ctx, - tx, - agentInput, - contentBlocks, - ); err != nil { + parts, err := executionContent(contentBlocks) + if err != nil { return createAgentContentInputTxResult{}, err } - result := createAgentContentInputTxResult{ - agentInput: agentInput, - contentBlocks: input.ContentBlocks, - created: true, - } - if input.CancelOpenInteractions { - result.canceledInteractionIDs, err = cancelOpenInteractionsForSteeringInputTx( - ctx, - txNotifications, - tx, - qtx, - input.ProjectID, - input.AgentID, - agentInput.ID, - ) - if err != nil { - return createAgentContentInputTxResult{}, err - } + received, err := h.ReceiveContent(ctx, agentexecution.ReceiveContentInput{ + ActorID: actorID, IntegrationTargetID: input.IntegrationTargetID, + DeliveryMode: string(input.DeliveryMode), IdempotencyScope: input.IdempotencyScope, + IdempotencyKey: input.IdempotencyKey, Metadata: input.Metadata, Content: parts, + }) + if err != nil { + return createAgentContentInputTxResult{}, err } - if err := qtx.ReconcileAgentWakeup( + row, err := qtx.GetAgentInput( ctx, - dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: []byte(`{"reason":"agent_input"}`), - }, - ); err != nil { - return createAgentContentInputTxResult{}, fmt.Errorf( - "reconcile agent wakeup after input: %w", - err, - ) + dbsqlc.GetAgentInputParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: received.ID}, + ) + if err != nil { + return createAgentContentInputTxResult{}, err } - return result, nil -} - -func createAgentInputContentBlocksTx( - ctx context.Context, - tx pgx.Tx, - input AgentInputRecord, - blocks []CreateContentBlockInput, -) error { - for _, block := range blocks { - block.ProjectID = input.ProjectID - block.AgentID = input.AgentID - block.OwnerKind = ContentBlockOwnerAgentInput - block.OwnerAgentInputID = input.ID - if _, err := createContentBlockTx(ctx, tx, block); err != nil { - return err - } + result := createAgentContentInputTxResult{agentInput: agentInputRecordFromGetSQLC(row), + contentBlocks: input.ContentBlocks, created: received.Created} + if received.Created && input.CancelOpenInteractions { + result.canceledInteractionIDs, err = h.SupersedeInteractions(ctx, received.ID) } - return nil + return result, err } func agentInputContentBlocks( diff --git a/internal/storage/executionstore/agent_fixture_integration.go b/internal/storage/executionstore/agent_fixture_integration.go index 188f7b0bf..e3efd1c61 100644 --- a/internal/storage/executionstore/agent_fixture_integration.go +++ b/internal/storage/executionstore/agent_fixture_integration.go @@ -9,8 +9,8 @@ import ( "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" ) @@ -90,13 +90,14 @@ func (s *Store) CreateAgentFixture(ctx context.Context, input AgentFixtureInput) if input.CurrentConfigID == uuid.Nil { return AgentRecord{}, errors.New("current config id is required") } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentRecord{}, fmt.Errorf("begin create agent fixture: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + qtx := dbsqlc.New(tx) project, err := loadProjectTx(ctx, qtx, input.ProjectID) if err != nil { return AgentRecord{}, err @@ -111,7 +112,10 @@ func (s *Store) CreateAgentFixture(ctx context.Context, input AgentFixtureInput) if err := lockAgentConfigForUseTx(ctx, qtx, config); err != nil { return AgentRecord{}, err } - record, _, err := insertAdmittedAgentTx(ctx, tx, qtx, insertAgentInput{ + if err := lockResourceCreation(ctx, qtx, resourceAgents, input.ProjectID.String()); err != nil { + return AgentRecord{}, err + } + record, _, err := insertAdmittedAgentTx(ctx, unit, qtx, insertAgentInput{ OrgID: project.OrgID, ProjectID: input.ProjectID, Name: input.Name, @@ -120,7 +124,7 @@ func (s *Store) CreateAgentFixture(ctx context.Context, input AgentFixtureInput) if err != nil { return AgentRecord{}, err } - if _, err := activateNewAgentConfigTx(ctx, txNotifications, tx, qtx, ActivateAgentConfigInput{ + if _, err := activateNewAgentConfigTx(ctx, unit, qtx, ActivateAgentConfigInput{ ProjectID: input.ProjectID, AgentID: record.ID, AgentConfigID: input.CurrentConfigID, @@ -130,7 +134,7 @@ func (s *Store) CreateAgentFixture(ctx context.Context, input AgentFixtureInput) }); err != nil { return AgentRecord{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "create agent fixture"); err != nil { + if err := unit.Commit(ctx, "create agent fixture"); err != nil { return AgentRecord{}, err } return record, nil diff --git a/internal/storage/executionstore/agent_input_admission_integration_test.go b/internal/storage/executionstore/agent_input_admission_integration_test.go index 45c2c592d..03e25c410 100644 --- a/internal/storage/executionstore/agent_input_admission_integration_test.go +++ b/internal/storage/executionstore/agent_input_admission_integration_test.go @@ -5,13 +5,10 @@ package executionstore_test import ( "context" "encoding/json" - "errors" "testing" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) func TestAgentInputAdmissionUsesItsCausalEvent(t *testing.T) { @@ -45,10 +42,18 @@ func TestAgentInputAdmissionUsesItsCausalEvent(t *testing.T) { admittedInput := admitted.Inputs[0] admittedEvent := admitted.Events[0] if admittedInput.AdmittedAt == nil || !admittedInput.AdmittedAt.Equal(admittedEvent.At) { - t.Fatalf("input admitted_at = %v, want event created_at %s", admittedInput.AdmittedAt, admittedEvent.At) + t.Fatalf( + "input admitted_at = %v, want event created_at %s", + admittedInput.AdmittedAt, + admittedEvent.At, + ) } if admittedInput.ResolvedAt == nil || !admittedInput.ResolvedAt.Equal(admittedEvent.At) { - t.Fatalf("input resolved_at = %v, want event created_at %s", admittedInput.ResolvedAt, admittedEvent.At) + t.Fatalf( + "input resolved_at = %v, want event created_at %s", + admittedInput.ResolvedAt, + admittedEvent.At, + ) } unrelatedInput, _, _, err := fixture.Store.Execution().CreateAgentContentInput( @@ -64,14 +69,14 @@ func TestAgentInputAdmissionUsesItsCausalEvent(t *testing.T) { if err != nil { t.Fatalf("create unrelated input: %v", err) } - _, err = fixture.Store.q.AdmitAgentInput(ctx, dbsqlc.AdmitAgentInputParams{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ID: unrelatedInput.ID, - AdmittedEventID: admittedEvent.ID, - }) - if !errors.Is(err, pgx.ErrNoRows) { - t.Fatalf("admit input with unrelated event error = %v, want %v", err, pgx.ErrNoRows) + _, err = fixture.Store.pool.Exec( + ctx, + `UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=statement_timestamp(),resolved_at=statement_timestamp() WHERE id=$1`, + unrelatedInput.ID, + admittedEvent.ID, + ) + if !isPgCheckViolation(err) { + t.Fatalf("admit unrelated event error=%v, want check violation", err) } var state string var admittedEventID *uuid.UUID @@ -85,6 +90,10 @@ WHERE project_id = $1 t.Fatalf("load unrelated input after rejected admission: %v", err) } if state != "received" || admittedEventID != nil { - t.Fatalf("unrelated input state=%q admitted_event_id=%v, want received with no event", state, admittedEventID) + t.Fatalf( + "unrelated input state=%q admitted_event_id=%v, want received with no event", + state, + admittedEventID, + ) } } diff --git a/internal/storage/executionstore/agent_input_admission_test_helpers_test.go b/internal/storage/executionstore/agent_input_admission_test_helpers_test.go index 41a697c7b..0dc54a8e5 100644 --- a/internal/storage/executionstore/agent_input_admission_test_helpers_test.go +++ b/internal/storage/executionstore/agent_input_admission_test_helpers_test.go @@ -8,9 +8,8 @@ import ( "testing" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) @@ -44,60 +43,38 @@ func admitNextAgentInputAndOpenTurnForTestErr( "project id, agent id, and runtime lock id are required", ) } - tx, err := store.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := store.Execution().IntegrationBeginUnit(ctx) if err != nil { - return executionstore.AdmittedAgentInputTurn{}, false, fmt.Errorf("begin admit next agent input: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject( - ctx, dbsqlc.LockAgentInProjectParams{ProjectID: projectID, ID: agentID}, - ); err != nil { - return executionstore.AdmittedAgentInputTurn{}, false, fmt.Errorf("lock agent for input admission: %w", err) + return executionstore.AdmittedAgentInputTurn{}, false, fmt.Errorf( + "begin admit next agent input: %w", + err, + ) } - if err := executionstore.IntegrationEnsureRuntimeLockActiveTx(ctx, tx, projectID, agentID, runtimeLockID); err != nil { + defer func() { _ = unit.Rollback(ctx) }() + if _, + err := unit.LockAgent(ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: projectID, + ID: agentID}, + agentexecution.RuntimeAuthority{AgentID: agentID, + RuntimeLockID: runtimeLockID}); err != nil { return executionstore.AdmittedAgentInputTurn{}, false, err } - if incomplete, err := qtx.AgentHasIncompleteToolBatch( - ctx, - dbsqlc.AgentHasIncompleteToolBatchParams{ - ProjectID: projectID, - AgentID: agentID, - }, - ); err != nil { - return executionstore.AdmittedAgentInputTurn{}, false, fmt.Errorf("check incomplete tool batch: %w", err) - } else if incomplete { - return executionstore.AdmittedAgentInputTurn{}, false, nil - } - selected, err := executionstore.IntegrationSelectLockedSteeringAgentInputsForAdmissionTx(ctx, qtx, projectID, agentID) - if err != nil { + if err := executionstore.IntegrationEnsureRuntimeLockActiveTx(ctx, + unit, + projectID, + agentID, + runtimeLockID); err != nil { return executionstore.AdmittedAgentInputTurn{}, false, err } - if len(selected) == 0 { - selected, err = executionstore.IntegrationSelectLockedQueuedAgentInputForAdmissionTx(ctx, qtx, projectID, agentID) - if err != nil { - return executionstore.AdmittedAgentInputTurn{}, false, err - } - } - if len(selected) == 0 { - return executionstore.AdmittedAgentInputTurn{}, false, nil - } - admitted, err := executionstore.IntegrationAdmitLockedAgentInputsAndOpenTurnTx( - ctx, - notifications.NewTxNotifications(), - tx, - qtx, - executionstore.IntegrationAdmitAgentInputAndOpenTurnInput{ - ProjectID: projectID, - AgentID: agentID, - }, - selected, - ) + admitted, err := executionstore.IntegrationAdmitInputs(ctx, unit, projectID, agentID) if err != nil { return executionstore.AdmittedAgentInputTurn{}, false, err } - if err := tx.Commit(ctx); err != nil { - return executionstore.AdmittedAgentInputTurn{}, false, fmt.Errorf("commit admit next agent input: %w", err) + if err := unit.Commit(ctx, "input admission"); err != nil { + return executionstore.AdmittedAgentInputTurn{}, false, fmt.Errorf( + "commit admit next agent input: %w", + err, + ) } - return admitted, true, nil + return admitted, len(admitted.Inputs) > 0, nil } diff --git a/internal/storage/executionstore/agent_input_rows.go b/internal/storage/executionstore/agent_input_rows.go index 95f9069eb..be8a5f5dd 100644 --- a/internal/storage/executionstore/agent_input_rows.go +++ b/internal/storage/executionstore/agent_input_rows.go @@ -8,32 +8,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" ) -func agentInputRecordFromInsertSQLC(row dbsqlc.InsertAgentInputRow) AgentInputRecord { - return agentInputRecordFromNewFields( - row.ID, - row.ProjectID, - row.AgentID, - row.State, - row.InputRank, - row.ActorID, - row.InputKind, - row.IntegrationTargetID, - row.IdempotencyScope, - row.InputIdempotencyKey, - row.QueuedAt, - row.AdmittedEventID, - row.AdmittedAt, - row.CanceledAt, - row.DeliveryMode, - row.ControlType, - row.TargetInteractionID, - row.AgentConfigID, - row.ResolvedAt, - row.RejectedReason, - row.Metadata, - ) -} - func agentInputRecordFromIdempotencySQLC(row dbsqlc.GetAgentInputByIdempotencyRow) AgentInputRecord { return agentInputRecordFromNewFields( row.ID, @@ -60,112 +34,6 @@ func agentInputRecordFromIdempotencySQLC(row dbsqlc.GetAgentInputByIdempotencyRo ) } -func agentInputRecordFromControlSQLC(row dbsqlc.InsertControlAgentInputRow) AgentInputRecord { - return agentInputRecordFromNewFields( - row.ID, - row.ProjectID, - row.AgentID, - row.State, - row.InputRank, - row.ActorID, - row.InputKind, - nil, - row.IdempotencyScope, - row.InputIdempotencyKey, - row.QueuedAt, - row.AdmittedEventID, - row.AdmittedAt, - row.CanceledAt, - row.DeliveryMode, - row.ControlType, - row.TargetInteractionID, - row.AgentConfigID, - row.ResolvedAt, - row.RejectedReason, - row.Metadata, - ) -} - -func agentInputRecordFromInteractionResponseInsertSQLC( - row dbsqlc.InsertInteractionResponseAgentInputRow, -) AgentInputRecord { - return agentInputRecordFromNewFields( - row.ID, - row.ProjectID, - row.AgentID, - row.State, - row.InputRank, - row.ActorID, - row.InputKind, - nil, - row.IdempotencyScope, - row.InputIdempotencyKey, - row.QueuedAt, - row.AdmittedEventID, - row.AdmittedAt, - row.CanceledAt, - row.DeliveryMode, - row.ControlType, - row.TargetInteractionID, - row.AgentConfigID, - row.ResolvedAt, - row.RejectedReason, - row.Metadata, - ) -} - -func agentInputRecordFromSteeringAdmissionSQLC(row dbsqlc.ListSteeringAgentInputsForAdmissionRow) AgentInputRecord { - return agentInputRecordFromNewFields( - row.ID, - row.ProjectID, - row.AgentID, - row.State, - row.InputRank, - row.ActorID, - row.InputKind, - row.IntegrationTargetID, - row.IdempotencyScope, - row.InputIdempotencyKey, - row.QueuedAt, - row.AdmittedEventID, - row.AdmittedAt, - row.CanceledAt, - row.DeliveryMode, - row.ControlType, - row.TargetInteractionID, - nil, - row.ResolvedAt, - row.RejectedReason, - row.Metadata, - ) -} - -func agentInputRecordFromQueuedAdmissionSQLC(row dbsqlc.GetNextQueuedAgentInputForAdmissionRow) AgentInputRecord { - return agentInputRecordFromNewFields( - row.ID, - row.ProjectID, - row.AgentID, - row.State, - row.InputRank, - row.ActorID, - row.InputKind, - row.IntegrationTargetID, - row.IdempotencyScope, - row.InputIdempotencyKey, - row.QueuedAt, - row.AdmittedEventID, - row.AdmittedAt, - row.CanceledAt, - row.DeliveryMode, - row.ControlType, - row.TargetInteractionID, - nil, - row.ResolvedAt, - row.RejectedReason, - row.Metadata, - ) -} - func agentInputRecordFromBacklogSQLC(row dbsqlc.ListQueuedBacklogInputsRow) AgentInputRecord { return agentInputRecordFromNewFields( row.ID, @@ -237,3 +105,29 @@ func agentInputRecordFromNewFields( Metadata: metadata, } } + +func agentInputRecordFromGetSQLC(row dbsqlc.GetAgentInputRow) AgentInputRecord { + return agentInputRecordFromNewFields( + row.ID, + row.ProjectID, + row.AgentID, + row.State, + row.InputRank, + row.ActorID, + row.InputKind, + row.IntegrationTargetID, + row.IdempotencyScope, + row.InputIdempotencyKey, + row.QueuedAt, + row.AdmittedEventID, + row.AdmittedAt, + row.CanceledAt, + row.DeliveryMode, + row.ControlType, + row.TargetInteractionID, + row.AgentConfigID, + row.ResolvedAt, + row.RejectedReason, + row.Metadata, + ) +} diff --git a/internal/storage/executionstore/agent_inputs_backlog_integration_test.go b/internal/storage/executionstore/agent_inputs_backlog_integration_test.go index 0c7f394d8..ffb571348 100644 --- a/internal/storage/executionstore/agent_inputs_backlog_integration_test.go +++ b/internal/storage/executionstore/agent_inputs_backlog_integration_test.go @@ -11,7 +11,6 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -629,17 +628,17 @@ func TestCreateAgentContentInputOrdersQueueTimeAfterAgentLock(t *testing.T) { agentID := mustCreateAgent(t, ctx, store) actor := mustOmnaraActorParams(t, user.ID) - waitingTx, err := pool.Begin(ctx) + waitingTx, err := store.Execution().IntegrationBeginUnit(ctx) if err != nil { t.Fatalf("begin waiting input transaction: %v", err) } defer func() { _ = waitingTx.Rollback(ctx) }() - blockingTx, err := pool.Begin(ctx) + blockingTx, err := store.Execution().IntegrationBeginUnit(ctx) if err != nil { t.Fatalf("begin blocking input transaction: %v", err) } defer func() { _ = blockingTx.Rollback(ctx) }() - blockingQueries := store.q.WithTx(blockingTx) + blockingQueries := dbsqlc.New(blockingTx.DB()) if _, err := blockingQueries.LockAgentInProject( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: agentID}, @@ -647,7 +646,7 @@ func TestCreateAgentContentInputOrdersQueueTimeAfterAgentLock(t *testing.T) { t.Fatalf("lock agent for first input: %v", err) } var blockingPID int32 - if err := blockingTx.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&blockingPID); err != nil { + if err := blockingTx.DB().QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&blockingPID); err != nil { t.Fatalf("get blocking input backend: %v", err) } @@ -664,9 +663,8 @@ func TestCreateAgentContentInputOrdersQueueTimeAfterAgentLock(t *testing.T) { go func() { result, err := executionstore.IntegrationCreateAgentContentInputTx( context.Background(), - notifications.NewTxNotifications(), waitingTx, - store.q.WithTx(waitingTx), + dbsqlc.New(waitingTx.DB()), executionstore.CreateAgentContentInputInput{ ProjectID: testProjectID, AgentID: agentID, @@ -677,7 +675,7 @@ func TestCreateAgentContentInputOrdersQueueTimeAfterAgentLock(t *testing.T) { secondBlocks, ) if err == nil { - err = waitingTx.Commit(context.Background()) + err = waitingTx.Commit(context.Background(), "second input") } waitingResult <- createResult{input: result.AgentInput, err: err} }() @@ -690,7 +688,6 @@ func TestCreateAgentContentInputOrdersQueueTimeAfterAgentLock(t *testing.T) { } first, err := executionstore.IntegrationCreateAgentContentInputTx( ctx, - notifications.NewTxNotifications(), blockingTx, blockingQueries, executionstore.CreateAgentContentInputInput{ @@ -706,10 +703,10 @@ func TestCreateAgentContentInputOrdersQueueTimeAfterAgentLock(t *testing.T) { t.Fatalf("create first input: %v", err) } var releaseFloor time.Time - if err := blockingTx.QueryRow(ctx, `SELECT statement_timestamp()`).Scan(&releaseFloor); err != nil { + if err := blockingTx.DB().QueryRow(ctx, `SELECT statement_timestamp()`).Scan(&releaseFloor); err != nil { t.Fatalf("read input lock release floor: %v", err) } - if err := blockingTx.Commit(ctx); err != nil { + if err := blockingTx.Commit(ctx, "first input"); err != nil { t.Fatalf("commit first input: %v", err) } diff --git a/internal/storage/executionstore/agent_inputs_backlog_store.go b/internal/storage/executionstore/agent_inputs_backlog_store.go index d9bc82cc4..17b34800e 100644 --- a/internal/storage/executionstore/agent_inputs_backlog_store.go +++ b/internal/storage/executionstore/agent_inputs_backlog_store.go @@ -2,14 +2,10 @@ package executionstore import ( "context" - "encoding/json" "errors" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) type PromoteQueuedInputToSteeringInput struct { @@ -29,24 +25,15 @@ func (s *Store) PromoteQueuedInputToSteering( ctx context.Context, input PromoteQueuedInputToSteeringInput, ) error { - return s.changeAgentInputDeliveryMode( + return s.changeBacklog( ctx, input.ProjectID, input.AgentID, - input.InputID, - input.CancelOpenInteractions, - func(qtx *dbsqlc.Queries, ctx context.Context) (bool, bool, error) { - result, err := qtx.PromoteQueuedInputToSteering( - ctx, - dbsqlc.PromoteQueuedInputToSteeringParams{ - RankStride: agentInputRankStride, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - }, - ) - return result.Changed, result.Effective != nil && *result.Effective, err + agentexecution.BacklogChange{ + ID: input.InputID, + DeliveryMode: "steering", }, + input.CancelOpenInteractions, ) } @@ -54,83 +41,28 @@ func (s *Store) DemoteSteeringInputToQueued( ctx context.Context, input DemoteSteeringInputToQueuedInput, ) error { - return s.changeAgentInputDeliveryMode( - ctx, - input.ProjectID, - input.AgentID, - input.InputID, - false, - func(qtx *dbsqlc.Queries, ctx context.Context) (bool, bool, error) { - changed, err := qtx.DemoteSteeringInputToQueued( - ctx, - dbsqlc.DemoteSteeringInputToQueuedParams{ - RankStride: agentInputRankStride, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - }, - ) - return changed == 1, changed == 1, err - }, - ) + return s.changeBacklog(ctx, input.ProjectID, input.AgentID, + agentexecution.BacklogChange{ID: input.InputID, DeliveryMode: "queued"}, false) } -func (s *Store) changeAgentInputDeliveryMode( - ctx context.Context, - projectID, agentID, inputID uuid.UUID, - cancelOpenInteractions bool, - mutate func(*dbsqlc.Queries, context.Context) (bool, bool, error), -) error { - if projectID == uuid.Nil || agentID == uuid.Nil || inputID == uuid.Nil { - return errors.New("project id, agent id, and input id are required") +func (s *Store) changeBacklog(ctx context.Context, projectID, agentID uuid.UUID, + change agentexecution.BacklogChange, cancelInteractions bool) error { + if projectID == uuid.Nil || agentID == uuid.Nil || change.ID == uuid.Nil { + return errors.New("project, agent and input are required") } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, projectID, agentID) if err != nil { - return fmt.Errorf("begin change agent input delivery mode: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: projectID, ID: agentID}, - ); err != nil { - return fmt.Errorf("lock agent for input delivery mode change: %w", err) + return err } - changed, effective, err := mutate(qtx, ctx) + defer func() { _ = unit.Rollback(ctx) }() + changed, err := h.ChangeBacklog(ctx, change) if err != nil { - return fmt.Errorf("change agent input delivery mode: %w", err) - } - if !effective { - return storeerr.ErrStateTransitionConflict + return err } - if !changed { - return nil - } - if cancelOpenInteractions { - if _, err := cancelOpenInteractionsForSteeringInputTx( - ctx, - txNotifications, - tx, - qtx, - projectID, - agentID, - inputID, - ); err != nil { + if changed && cancelInteractions { + if _, err = h.SupersedeInteractions(ctx, change.ID); err != nil { return err } } - if err := qtx.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - Metadata: json.RawMessage(`{"reason":"input_delivery_mode_changed"}`), - }); err != nil { - return fmt.Errorf("reconcile wakeup after input delivery mode change: %w", err) - } - return s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "change agent input delivery mode", - ) + return unit.Commit(ctx, "change input backlog") } diff --git a/internal/storage/executionstore/agent_inputs_ingress_store.go b/internal/storage/executionstore/agent_inputs_ingress_store.go index fa770cc92..ca2838d7e 100644 --- a/internal/storage/executionstore/agent_inputs_ingress_store.go +++ b/internal/storage/executionstore/agent_inputs_ingress_store.go @@ -10,8 +10,6 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" ) type AgentInputRecord struct { @@ -47,8 +45,6 @@ type AgentInputQueueCursor struct { ID uuid.UUID } -const agentInputRankStride int64 = 1024 - type ListQueuedBacklogInputsInput struct { ProjectID uuid.UUID AgentID uuid.UUID @@ -69,57 +65,9 @@ const ( type AgentInputDeliveryMode string -type insertAgentInputInput struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - DeliveryMode AgentInputDeliveryMode - ActorID uuid.UUID - IntegrationTargetID uuid.UUID - IdempotencyScope string - InputIdempotencyKey string - Metadata json.RawMessage -} - -func insertAgentInputTx( - ctx context.Context, - tx pgx.Tx, - input insertAgentInputInput, -) (AgentInputRecord, error) { - if input.DeliveryMode == "" { - input.DeliveryMode = DeliveryModeQueued - } - if input.DeliveryMode != DeliveryModeQueued && input.DeliveryMode != DeliveryModeSteering { - return AgentInputRecord{}, fmt.Errorf( - "unsupported agent input delivery_mode %q", - input.DeliveryMode, - ) - } - input.Metadata = normalizedJSON(input.Metadata) - row, err := dbsqlc.New(tx).InsertAgentInput(ctx, dbsqlc.InsertAgentInputParams{ - RankStride: agentInputRankStride, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: storeutil.IDFromNil(input.ID), - DeliveryMode: string(input.DeliveryMode), - ActorID: storeutil.IDFromNil(input.ActorID), - IntegrationTargetID: storeutil.IDFromNil(input.IntegrationTargetID), - IdempotencyScope: storeutil.TextFromEmpty(input.IdempotencyScope), - InputIdempotencyKey: storeutil.TextFromEmpty(input.InputIdempotencyKey), - Metadata: input.Metadata, - }) - if err != nil { - if storeutil.IsUniqueViolation(err) { - return AgentInputRecord{}, storeerr.ErrIdempotencyConflict - } - return AgentInputRecord{}, fmt.Errorf("insert agent input: %w", err) - } - return agentInputRecordFromInsertSQLC(row), nil -} - func loadAgentInputByIdempotencyMaybeTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID uuid.UUID, scope, key string, ) (AgentInputRecord, bool, error) { diff --git a/internal/storage/executionstore/agent_inputs_store.go b/internal/storage/executionstore/agent_inputs_store.go index 84ffbbf0e..847b1a41b 100644 --- a/internal/storage/executionstore/agent_inputs_store.go +++ b/internal/storage/executionstore/agent_inputs_store.go @@ -9,7 +9,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/events" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -59,6 +59,7 @@ const ( ) type ClaimedModelWork struct { + Prepared *PreparedNormalModelCall Kind ModelWorkKind ModelCallContextID uuid.UUID SourceModelCallContextID uuid.UUID @@ -70,6 +71,7 @@ type ClaimedModelWork struct { } type ClaimedToolWork struct { + Prepared *PreparedToolWork TurnID uuid.UUID ModelCallContextID uuid.UUID ModelOutputID uuid.UUID @@ -77,18 +79,14 @@ type ClaimedToolWork struct { } type ClaimedAgentWork struct { - OrgID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - Kind AgentWorkKind - RuntimeLock AgentRuntimeLockRecord - Model ClaimedModelWork - Tool ClaimedToolWork -} - -type admitAgentInputAndOpenTurnInput struct { - ProjectID uuid.UUID - AgentID uuid.UUID + LocalLeaseBudgetStartedAt time.Time + OrgID uuid.UUID + ProjectID uuid.UUID + AgentID uuid.UUID + Kind AgentWorkKind + RuntimeLock AgentRuntimeLockRecord + Model ClaimedModelWork + Tool ClaimedToolWork } type AdmittedAgentInputTurn struct { @@ -97,347 +95,49 @@ type AdmittedAgentInputTurn struct { Turn AgentTurnRecord } -func (s *Store) ClaimNextAgentWork(ctx context.Context, input ClaimNextAgentWorkInput) (ClaimedAgentWork, bool, error) { +func (s *Store) ClaimNextAgentWork( + ctx context.Context, + input ClaimNextAgentWorkInput, +) (ClaimedAgentWork, bool, error) { if input.WorkerProcessID == uuid.Nil { return ClaimedAgentWork{}, false, errors.New("worker process id is required") } if err := validateAgentRuntimeLockLeaseDuration(input.LeaseDuration); err != nil { return ClaimedAgentWork{}, false, err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { - return ClaimedAgentWork{}, false, fmt.Errorf("begin claim agent work: %w", err) + return ClaimedAgentWork{}, false, err } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - wakeup, err := qtx.ClaimNextAgentWakeup(ctx) + defer func() { _ = unit.Rollback(ctx) }() + route, err := unit.ClaimAgent(ctx) if errors.Is(err, pgx.ErrNoRows) { return ClaimedAgentWork{}, false, storeerr.ErrNoClaimableAgentWakeup } if err != nil { - return ClaimedAgentWork{}, false, fmt.Errorf("claim agent wakeup: %w", err) - } - claim := ClaimedAgentWork{ProjectID: wakeup.ProjectID, AgentID: wakeup.AgentID} - - toolWork, toolErr := qtx.NextAgentToolWork( - ctx, - dbsqlc.NextAgentToolWorkParams{ProjectID: wakeup.ProjectID, AgentID: wakeup.AgentID}, - ) - hasToolWork := toolErr == nil - if toolErr != nil && !errors.Is(toolErr, pgx.ErrNoRows) { - return ClaimedAgentWork{}, false, fmt.Errorf("load next tool work: %w", toolErr) - } - modelWork, modelErr := qtx.NextAgentModelWork(ctx, dbsqlc.NextAgentModelWorkParams{ - ProjectID: wakeup.ProjectID, - AgentID: wakeup.AgentID, - }) - hasModelWork := modelErr == nil - if modelErr != nil && !errors.Is(modelErr, pgx.ErrNoRows) { - return ClaimedAgentWork{}, false, fmt.Errorf("load next model work: %w", modelErr) - } - if hasToolWork && hasModelWork { - return ClaimedAgentWork{}, false, errors.New( - "agent has both tool work and model work for its current frontier", - ) - } - if hasToolWork { - claim.Kind = AgentWorkTool - claim.Tool = ClaimedToolWork{ - TurnID: toolWork.TurnID, - ModelCallContextID: toolWork.ModelCallContextID, - ModelOutputID: toolWork.ModelOutputID, - SourceEventID: toolWork.SourceEventID, - } - } - - var selectedInputs []AgentInputRecord - if claim.Kind != AgentWorkTool { - incompleteToolBatch, err := qtx.AgentHasIncompleteToolBatch( - ctx, - dbsqlc.AgentHasIncompleteToolBatchParams{ - ProjectID: wakeup.ProjectID, - AgentID: wakeup.AgentID, - }, - ) - if err != nil { - return ClaimedAgentWork{}, false, fmt.Errorf("check incomplete tool batch: %w", err) - } - if incompleteToolBatch { - if hasModelWork { - return ClaimedAgentWork{}, false, errors.New( - "agent has model work while its current tool batch is incomplete", - ) - } - return s.reconcileClaimedWakeupTx( - ctx, - txNotifications, - tx, - qtx, - claim, - "tool_batch_wait", - ) - } - - steering, err := selectLockedSteeringAgentInputsForAdmissionTx( - ctx, - qtx, - wakeup.ProjectID, - wakeup.AgentID, - ) - if err != nil { - return ClaimedAgentWork{}, false, err - } - switch { - case len(steering) > 0: - claim.Kind = AgentWorkModel - claim.Model.Kind = ModelWorkStart - selectedInputs = steering - case hasModelWork && modelWork.IsReady: - claim.Kind = AgentWorkModel - claim.Model = claimedModelWorkFromSQLC(modelWork) - case hasModelWork: - return s.reconcileClaimedWakeupTx( - ctx, - txNotifications, - tx, - qtx, - claim, - "model_retry_wait", - ) - default: - queued, err := selectLockedQueuedAgentInputForAdmissionTx( - ctx, - qtx, - wakeup.ProjectID, - wakeup.AgentID, - ) - if err != nil { - return ClaimedAgentWork{}, false, err - } - if len(queued) == 0 { - return s.reconcileClaimedWakeupTx( - ctx, - txNotifications, - tx, - qtx, - claim, - "idle", - ) - } - claim.Kind = AgentWorkModel - claim.Model.Kind = ModelWorkStart - selectedInputs = queued - } - } - - runtime, orgID, err := acquireAgentRuntimeLockTx( - ctx, - qtx, - wakeup.ProjectID, - wakeup.AgentID, - input.WorkerProcessID, - input.LeaseDuration, - ) - if err != nil { - if errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - return ClaimedAgentWork{}, true, nil - } return ClaimedAgentWork{}, false, err } - claim.RuntimeLock = runtime - claim.OrgID = orgID - if err := consumeClaimedAgentWakeupTx( - ctx, - qtx, - claim.ProjectID, - claim.AgentID, - ); err != nil { + h, err := unit.Agent(route) + if err != nil { return ClaimedAgentWork{}, false, err } - - if len(selectedInputs) > 0 { - admitted, err := admitLockedAgentInputsAndOpenTurnTx( - ctx, - txNotifications, - tx, - qtx, - admitAgentInputAndOpenTurnInput{ - ProjectID: wakeup.ProjectID, - AgentID: wakeup.AgentID, - }, - selectedInputs, - ) - if err != nil { - return ClaimedAgentWork{}, false, err - } - claim.Model.AdmittedInputTurn = admitted - claim.Model.TurnID = admitted.Turn.ID - claim.Model.InputIDs, claim.Model.OpeningEventSequence, err = modelCallOpeningInputSet( - ctx, - qtx, - claim.ProjectID, - claim.AgentID, - claim.Model.TurnID, - latestEventSequence(admitted.Events), - ) - if err != nil { - return ClaimedAgentWork{}, false, err - } - } - if claim.Kind == AgentWorkModel { - if err := validateClaimedModelWork(claim.Model); err != nil { - return ClaimedAgentWork{}, false, err - } - } - - renewal, err := renewAgentRuntimeLockTx( - ctx, - qtx, - claim.ProjectID, - claim.AgentID, - claim.RuntimeLock.ID, - input.LeaseDuration, - ) + started := time.Now() //nolint:omnaralint // Monotonic lease budget. + work, err := h.Claim(ctx, input.WorkerProcessID, input.LeaseDuration) if err != nil { return ClaimedAgentWork{}, false, err } - claim.RuntimeLock = renewal.RuntimeLock - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "claim agent work"); err != nil { + if err := applyAdmissionDestination(ctx, unit, route.ProjectID, route.AgentID, work.Admission); err != nil { return ClaimedAgentWork{}, false, err } - return claim, true, nil -} - -func claimedModelWorkFromSQLC(row dbsqlc.NextAgentModelWorkRow) ClaimedModelWork { - work := ClaimedModelWork{ - Kind: ModelWorkKind(row.WorkKind), - TurnID: row.TurnID, - InputIDs: row.InputIds, - OpeningEventSequence: row.OpeningEventSequence, - } - switch work.Kind { - case ModelWorkStart: - case ModelWorkResume: - work.ModelCallContextID = row.ModelCallContextID - case ModelWorkContinue: - work.SourceModelCallContextID = row.ModelCallContextID - work.SourceModelOutputID = row.ModelOutputID - } - return work -} - -func (s *Store) reconcileClaimedWakeupTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - claim ClaimedAgentWork, - reason string, -) (ClaimedAgentWork, bool, error) { - metadata, err := marshalJSON(map[string]any{"reason": reason}) + claim, err := shapeOwnedWork(ctx, unit, route.ProjectID, route.AgentID, work) if err != nil { - return ClaimedAgentWork{}, false, fmt.Errorf("marshal wakeup reconciliation metadata: %w", err) - } - if err := qtx.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: claim.ProjectID, - AgentID: claim.AgentID, - Metadata: metadata, - }); err != nil { - return ClaimedAgentWork{}, false, fmt.Errorf("reconcile claimed agent wakeup: %w", err) - } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "reconcile claimed agent wakeup"); err != nil { return ClaimedAgentWork{}, false, err } - return ClaimedAgentWork{}, true, nil -} - -func consumeClaimedAgentWakeupTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) error { - changed, err := qtx.ConsumeAgentWakeup( - ctx, - dbsqlc.ConsumeAgentWakeupParams{ProjectID: projectID, AgentID: agentID}, - ) - if err != nil { - return fmt.Errorf("consume claimed agent wakeup: %w", err) - } - if changed != 1 { - return storeerr.ErrNoClaimableAgentWakeup - } - return nil -} - -func validateClaimedModelWork(work ClaimedModelWork) error { - if work.TurnID == uuid.Nil || len(work.InputIDs) == 0 || work.OpeningEventSequence <= 0 { - return errors.New("claimed model work requires a turn, opening inputs, and opening event sequence") - } - switch work.Kind { - case ModelWorkStart: - if work.ModelCallContextID != uuid.Nil || - work.SourceModelCallContextID != uuid.Nil || - work.SourceModelOutputID != uuid.Nil { - return errors.New("start model work has invalid source identity") - } - case ModelWorkResume: - if work.ModelCallContextID == uuid.Nil || - work.SourceModelCallContextID != uuid.Nil || - work.SourceModelOutputID != uuid.Nil { - return errors.New("resume model work requires only its active context") - } - case ModelWorkContinue: - if work.ModelCallContextID != uuid.Nil || - work.SourceModelCallContextID == uuid.Nil || - work.SourceModelOutputID == uuid.Nil { - return errors.New("continue model work requires its source context and source output") - } - default: - return fmt.Errorf("unsupported model work kind %q", work.Kind) - } - return nil -} - -func latestEventSequence(events []events.Event) int64 { - var sequence int64 - for _, event := range events { - if event.Sequence > sequence { - sequence = event.Sequence - } - } - return sequence -} - -func modelCallOpeningInputSet( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, turnID uuid.UUID, - inputEventSequence int64, -) ([]uuid.UUID, int64, error) { - rows, err := qtx.ListModelCallOpeningContentInputs( - ctx, - dbsqlc.ListModelCallOpeningContentInputsParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: turnID, - InputEventSequence: inputEventSequence, - }, - ) - if err != nil { - return nil, 0, fmt.Errorf("list model call opening inputs: %w", err) - } - if len(rows) == 0 { - return nil, 0, fmt.Errorf( - "model call opening inputs are empty: %w", - storeerr.ErrStateTransitionConflict, - ) - } - inputIDs := make([]uuid.UUID, 0, len(rows)) - for _, row := range rows { - inputIDs = append(inputIDs, row.InputID) + claim.LocalLeaseBudgetStartedAt = started + if err = unit.Commit(ctx, "claim agent work"); err != nil { + return ClaimedAgentWork{}, false, err } - return inputIDs, rows[0].EventSequence, nil + return claim, true, nil } func (s *Store) ListQueuedBacklogInputs( @@ -491,311 +191,22 @@ func (s *Store) ListQueuedBacklogInputs( return result, nil } -func selectLockedSteeringAgentInputsForAdmissionTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) ([]AgentInputRecord, error) { - steeringRows, err := qtx.ListSteeringAgentInputsForAdmission( - ctx, - dbsqlc.ListSteeringAgentInputsForAdmissionParams{ProjectID: projectID, AgentID: agentID}, - ) - if err != nil { - return nil, fmt.Errorf("list steering agent inputs: %w", err) - } - selected := make([]AgentInputRecord, 0, len(steeringRows)) - for _, row := range steeringRows { - selected = append(selected, agentInputRecordFromSteeringAdmissionSQLC(row)) - } - return selected, nil +func (s *Store) CancelQueuedBacklogInput(ctx context.Context, input CancelQueuedBacklogInputInput) error { + return s.changeBacklog(ctx, input.ProjectID, input.AgentID, + agentexecution.BacklogChange{ID: input.InputID, DeliveryMode: "queued", Cancel: true}, false) } -func selectLockedQueuedAgentInputForAdmissionTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) ([]AgentInputRecord, error) { - row, err := qtx.GetNextQueuedAgentInputForAdmission( +func (s *Store) MoveQueuedBacklogInput(ctx context.Context, input MoveQueuedBacklogInputInput) error { + return s.changeBacklog( ctx, - dbsqlc.GetNextQueuedAgentInputForAdmissionParams{ProjectID: projectID, AgentID: agentID}, - ) - if errors.Is(err, pgx.ErrNoRows) { - return nil, nil - } - if err != nil { - return nil, fmt.Errorf("get next queued agent input: %w", err) - } - return []AgentInputRecord{agentInputRecordFromQueuedAdmissionSQLC(row)}, nil -} - -func admitLockedAgentInputsAndOpenTurnTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - input admitAgentInputAndOpenTurnInput, - lockedInputs []AgentInputRecord, -) (AdmittedAgentInputTurn, error) { - if len(lockedInputs) == 0 { - return AdmittedAgentInputTurn{}, storeerr.ErrStateTransitionConflict - } - if len(lockedInputs) > 1 { - for _, locked := range lockedInputs { - if locked.DeliveryMode != DeliveryModeSteering { - return AdmittedAgentInputTurn{}, storeerr.ErrStateTransitionConflict - } - } - } - admittedInputs := make([]AgentInputRecord, 0, len(lockedInputs)) - admittedEvents := make([]events.Event, 0, len(lockedInputs)) - turnID, err := uuid.NewV7() - if err != nil { - return AdmittedAgentInputTurn{}, fmt.Errorf("generate turn id: %w", err) - } - for _, agentInput := range lockedInputs { - eventRecord, err := appendTypedAgentEventTx( - ctx, - txNotifications, - tx, - AppendTypedAgentEventInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: turnID, - IsOpeningEvent: true, - Kind: events.KindAgentInput, - IdempotencyKey: "agent_input:" + agentInput.ID.String(), - AgentInputID: agentInput.ID, - }, - ) - if err != nil { - return AdmittedAgentInputTurn{}, err - } - event := eventRecord.Event - admission, err := qtx.AdmitAgentInput( - ctx, - dbsqlc.AdmitAgentInputParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: agentInput.ID, - AdmittedEventID: event.ID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return AdmittedAgentInputTurn{}, storeerr.ErrStateTransitionConflict - } - if err != nil { - return AdmittedAgentInputTurn{}, fmt.Errorf("admit agent input: %w", err) - } - agentInput.State = "resolved" - agentInput.AdmittedEventID = event.ID - agentInput.AdmittedAt = admission.AdmittedAt - agentInput.ResolvedAt = admission.ResolvedAt - admittedInputs = append(admittedInputs, agentInput) - admittedEvents = append(admittedEvents, event) - } - if err := selectAdmittedInteractionDestinationTx(ctx, tx, input.ProjectID, input.AgentID, admittedInputs); err != nil { - return AdmittedAgentInputTurn{}, fmt.Errorf("select admitted input interaction destination: %w", err) - } - sequence, err := qtx.NextTurnSequence( - ctx, - dbsqlc.NextTurnSequenceParams{ProjectID: input.ProjectID, AgentID: input.AgentID}, - ) - if err != nil { - return AdmittedAgentInputTurn{}, fmt.Errorf("next turn sequence: %w", err) - } - latestEvent := admittedEvents[len(admittedEvents)-1] - turnRow, err := qtx.InsertAgentTurn( - ctx, - dbsqlc.InsertAgentTurnParams{ - ID: turnID, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnSequence: sequence, - LatestEventID: latestEvent.ID, - LatestSemanticEventID: latestEvent.ID, + input.ProjectID, + input.AgentID, + agentexecution.BacklogChange{ + ID: input.InputID, + DeliveryMode: "queued", + Position: string(input.Position), + BeforeID: input.AnchorInputID, }, + false, ) - if err != nil { - return AdmittedAgentInputTurn{}, fmt.Errorf("insert agent turn: %w", err) - } - return AdmittedAgentInputTurn{ - Inputs: admittedInputs, - Events: admittedEvents, - Turn: agentTurnRecordFromInsertSQLC(turnRow), - }, nil -} - -func (s *Store) CancelQueuedBacklogInput( - ctx context.Context, - input CancelQueuedBacklogInputInput, -) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.InputID == uuid.Nil { - return errors.New("project id, agent id, and input id are required") - } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return fmt.Errorf("begin cancel queued backlog input: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - ); err != nil { - return fmt.Errorf("lock agent for backlog cancel: %w", err) - } - changed, err := qtx.CancelQueuedBacklogInput( - ctx, - dbsqlc.CancelQueuedBacklogInputParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - }, - ) - if err != nil { - return fmt.Errorf("cancel queued backlog input: %w", err) - } - if changed != 1 { - return storeerr.ErrStateTransitionConflict - } - if err := qtx.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: []byte(`{"reason":"queued_input_canceled"}`), - }); err != nil { - return fmt.Errorf("reconcile wakeup after queued backlog cancel: %w", err) - } - if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("commit cancel queued backlog input: %w", err) - } - return nil -} - -func (s *Store) MoveQueuedBacklogInput( - ctx context.Context, - input MoveQueuedBacklogInputInput, -) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.InputID == uuid.Nil { - return errors.New("project id, agent id, and input id are required") - } - switch input.Position { - case MoveQueuedBacklogInputToFront, MoveQueuedBacklogInputToBack: - if input.AnchorInputID != uuid.Nil { - return errors.New("front/back moves do not accept an anchor input id") - } - case MoveQueuedBacklogInputBefore, MoveQueuedBacklogInputAfter: - if input.AnchorInputID == uuid.Nil { - return errors.New("before/after moves require an anchor input id") - } - if input.AnchorInputID == input.InputID { - return storeerr.ErrStateTransitionConflict - } - default: - return errors.New("position must be front, back, before, or after") - } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return fmt.Errorf("begin move queued backlog input: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - ); err != nil { - return fmt.Errorf("lock agent for backlog move: %w", err) - } - valid, err := qtx.QueuedBacklogMoveIsValid(ctx, dbsqlc.QueuedBacklogMoveIsValidParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - RequiresAnchor: input.Position == MoveQueuedBacklogInputBefore || - input.Position == MoveQueuedBacklogInputAfter, - AnchorID: input.AnchorInputID, - }) - if err != nil { - return fmt.Errorf("validate queued backlog move: %w", err) - } - if !valid { - return storeerr.ErrStateTransitionConflict - } - changed, err := moveQueuedBacklogInputOnce(ctx, qtx, input) - if err != nil { - return fmt.Errorf("move queued backlog input: %w", err) - } - if changed == 0 { - if _, err := qtx.RebalanceQueuedBacklogRanks( - ctx, - dbsqlc.RebalanceQueuedBacklogRanksParams{ - RankStride: agentInputRankStride, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }, - ); err != nil { - return fmt.Errorf("rebalance queued backlog ranks: %w", err) - } - changed, err = moveQueuedBacklogInputOnce(ctx, qtx, input) - if err != nil { - return fmt.Errorf("move queued backlog input after rank rebalance: %w", err) - } - } - if changed != 1 { - return storeerr.ErrStateTransitionConflict - } - if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("commit move queued backlog input: %w", err) - } - return nil -} - -func moveQueuedBacklogInputOnce( - ctx context.Context, - qtx *dbsqlc.Queries, - input MoveQueuedBacklogInputInput, -) (int64, error) { - switch input.Position { - case MoveQueuedBacklogInputToFront: - return qtx.MoveQueuedBacklogInputToFront( - ctx, - dbsqlc.MoveQueuedBacklogInputToFrontParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - RankStride: agentInputRankStride, - }, - ) - case MoveQueuedBacklogInputToBack: - return qtx.MoveQueuedBacklogInputToBack( - ctx, - dbsqlc.MoveQueuedBacklogInputToBackParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - RankStride: agentInputRankStride, - }, - ) - case MoveQueuedBacklogInputBefore: - return qtx.MoveQueuedBacklogInputBefore( - ctx, - dbsqlc.MoveQueuedBacklogInputBeforeParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - AnchorID: input.AnchorInputID, - }, - ) - case MoveQueuedBacklogInputAfter: - return qtx.MoveQueuedBacklogInputAfter( - ctx, - dbsqlc.MoveQueuedBacklogInputAfterParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.InputID, - AnchorID: input.AnchorInputID, - RankStride: agentInputRankStride, - }, - ) - default: - return 0, errors.New("invalid queued backlog move position") - } } diff --git a/internal/storage/executionstore/agent_integration_targets.go b/internal/storage/executionstore/agent_integration_targets.go index 5c6dd30ce..c0337013d 100644 --- a/internal/storage/executionstore/agent_integration_targets.go +++ b/internal/storage/executionstore/agent_integration_targets.go @@ -2,29 +2,17 @@ package executionstore import ( "context" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) type IntegrationAccess struct{} func (IntegrationAccess) ClearIntegrationTargetsFromAgents( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, projectID, integrationID uuid.UUID, ) error { - err := dbsqlc.New(tx).ClearDeletedIntegrationTargetsFromAgents( - ctx, - dbsqlc.ClearDeletedIntegrationTargetsFromAgentsParams{ - ProjectID: projectID, - IntegrationID: integrationID, - }, - ) - if err != nil { - return fmt.Errorf("clear integration targets from agents: %w", err) - } - return nil + return unit.ClearIntegrationTargets(ctx, projectID, integrationID) } diff --git a/internal/storage/executionstore/agent_interaction_cancellation_store.go b/internal/storage/executionstore/agent_interaction_cancellation_store.go deleted file mode 100644 index c39a31d02..000000000 --- a/internal/storage/executionstore/agent_interaction_cancellation_store.go +++ /dev/null @@ -1,80 +0,0 @@ -package executionstore - -import ( - "context" - "errors" - "fmt" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" -) - -const interactionCancellationReasonSupersededByInput = "superseded_by_input" - -func cancelOpenInteractionsForSteeringInputTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, - supersedingInputID uuid.UUID, -) ([]uuid.UUID, error) { - turn, err := qtx.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: projectID, AgentID: agentID}, - ) - if errors.Is(err, pgx.ErrNoRows) { - return nil, nil - } - if err != nil { - return nil, fmt.Errorf("load current turn for interaction cancellation: %w", err) - } - interactionIDs, err := qtx.CancelOpenAgentInteractionsForAgent( - ctx, - dbsqlc.CancelOpenAgentInteractionsForAgentParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: turn.ID, - Reason: interactionCancellationReasonSupersededByInput, - ResolvedByInputID: storeutil.IDFromNil(supersedingInputID), - }, - ) - if err != nil { - return nil, fmt.Errorf("cancel open interactions for steering input: %w", err) - } - rows, err := qtx.ListAgentInteractionsByIDs(ctx, dbsqlc.ListAgentInteractionsByIDsParams{ - ProjectID: projectID, - AgentID: agentID, - Ids: interactionIDs, - }) - if err != nil { - return nil, fmt.Errorf("load canceled interactions for steering input: %w", err) - } - canceledInteractionIDs := make([]uuid.UUID, 0, len(rows)) - for _, row := range rows { - interaction := agentInteractionRecordFromSQLC(row) - if err := applyPermissionInteractionResolutionTx( - ctx, - txNotifications, - tx, - qtx, - interaction, - ); err != nil { - return nil, err - } - if err := completeQuestionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - interaction, - ); err != nil { - return nil, err - } - canceledInteractionIDs = append(canceledInteractionIDs, interaction.ID) - } - return canceledInteractionIDs, nil -} diff --git a/internal/storage/executionstore/agent_interactions_store.go b/internal/storage/executionstore/agent_interactions_store.go index 1f57db858..f3be585cc 100644 --- a/internal/storage/executionstore/agent_interactions_store.go +++ b/internal/storage/executionstore/agent_interactions_store.go @@ -9,18 +9,14 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/interactionform" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/listing" - "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/toolpermission" ) -const interactionResponseIdempotencyScope = "agent_interaction_response" - type AgentInteractionState string type AgentInteractionKind string @@ -77,207 +73,79 @@ func (record AgentInteractionRecord) Form() (interactionform.Form, error) { return interactionFormForInteraction(record.InteractionKind, record.Request) } -func interactionUpdateFromRecord(record AgentInteractionRecord) *notifications.AgentInteractionUpdate { - return ¬ifications.AgentInteractionUpdate{ - ID: record.ID, InteractionKind: string(record.InteractionKind), State: string(record.State), - } -} - func (s *Store) CreatePermissionInteraction( ctx context.Context, input CreatePermissionInteractionInput, ) (AgentInteractionRecord, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ToolCallID == uuid.Nil || - input.RuntimeLockID == uuid.Nil { - return AgentInteractionRecord{}, errors.New( - "agent, tool call, and runtime lock are required", - ) - } - if err := input.Request.Validate(); err != nil { - return AgentInteractionRecord{}, fmt.Errorf("permission request: %w", err) - } - request, err := marshalJSON(input.Request) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("marshal permission request: %w", err) - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("begin create agent interaction: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return AgentInteractionRecord{}, err } - destination, err := captureInteractionDestinationTx(ctx, tx, input.ProjectID, input.AgentID) + defer func() { _ = unit.Rollback(ctx) }() + destination, err := captureInteractionDestinationTx(ctx, unit.DB(), input.ProjectID, input.AgentID) if err != nil { return AgentInteractionRecord{}, err } - interactionID, err := qtx.InsertAgentInteraction(ctx, dbsqlc.InsertAgentInteractionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ToolCallID: input.ToolCallID, - InteractionKind: string(AgentInteractionKindPermission), - Request: request, - Destination: sqlcRawMessageFromEmpty(destination), - }) + result, err := h.OpenInteraction(ctx, agentexecution.OpenInteractionInput{ToolRef: agentexecution.ToolRef{ + ID: input.ToolCallID, + RuntimeLockID: input.RuntimeLockID}, + Permission: &input.Request, + Destination: destination}) if err != nil { - if storeutil.IsUniqueViolation(err) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - return AgentInteractionRecord{}, fmt.Errorf("create agent interaction: %w", err) + return AgentInteractionRecord{}, err } - row, err := qtx.GetAgentInteraction(ctx, dbsqlc.GetAgentInteractionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: interactionID, - }) + row, err := dbsqlc.New(unit.DB()).GetAgentInteraction(ctx, dbsqlc.GetAgentInteractionParams{ + ProjectID: input.ProjectID, AgentID: input.AgentID, ID: result.ID}) if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("load created agent interaction: %w", err) - } - record := agentInteractionRecordFromSQLC(row) - if err := markToolCallAwaitingPermissionTx( - ctx, - txNotifications, - tx, - qtx, - input.ProjectID, - input.AgentID, - input.ToolCallID, - input.RuntimeLockID, - interactionUpdateFromRecord(record), - ); err != nil { - return AgentInteractionRecord{}, fmt.Errorf("link tool call to agent interaction: %w", err) + return AgentInteractionRecord{}, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "create agent interaction", - ); err != nil { + if err = unit.Commit(ctx, "open permission"); err != nil { return AgentInteractionRecord{}, err } - return record, nil + return agentInteractionRecordFromSQLC(row), nil } func (t *toolCallTransaction) createQuestionInteraction( ctx context.Context, input CreateQuestionInteractionInput, ) (AgentInteractionRecord, error) { - if err := input.Form.Validate(); err != nil { - return AgentInteractionRecord{}, fmt.Errorf("interaction form: %w", err) - } - request, err := marshalJSON(input.Form) - if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("marshal interaction form: %w", err) - } - existing, found, err := getAgentInteractionByToolCallKind( - ctx, - t.q, - t.input.ProjectID, - t.input.AgentID, - t.input.ToolCallID, - "question", - ) - if err != nil { - return AgentInteractionRecord{}, err - } - if found { - if !sameJSON(existing.Request, request) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - t.hasDurableCompletionOwner = true - if err := t.lockOrAcceptExisting(ctx); err != nil { - return AgentInteractionRecord{}, err - } - return existing, nil - } - if err := lockAgentWithParentTx(ctx, t.tx, t.q, t.input.ProjectID, t.input.AgentID); err != nil { - return AgentInteractionRecord{}, err - } - if err := t.lockForMutation(ctx); err != nil { + if err := lockAgentWithParentTx(ctx, t.unit, t.q, t.input.ProjectID, t.input.AgentID); err != nil { return AgentInteractionRecord{}, err } destination, err := captureInteractionDestinationTx(ctx, t.tx, t.input.ProjectID, t.input.AgentID) if err != nil { return AgentInteractionRecord{}, err } - interactionID, err := t.q.InsertAgentInteraction(ctx, dbsqlc.InsertAgentInteractionParams{ - ProjectID: t.input.ProjectID, - AgentID: t.input.AgentID, - ToolCallID: t.input.ToolCallID, - InteractionKind: string(AgentInteractionKindQuestion), - Request: request, - Destination: sqlcRawMessageFromEmpty(destination), - }) + h, err := t.unit.Handle(t.input.ProjectID, t.input.AgentID) if err != nil { - if storeutil.IsUniqueViolation(err) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - return AgentInteractionRecord{}, fmt.Errorf("create question interaction: %w", err) + return AgentInteractionRecord{}, err } - row, err := t.q.GetAgentInteraction(ctx, dbsqlc.GetAgentInteractionParams{ - ProjectID: t.input.ProjectID, - AgentID: t.input.AgentID, - ID: interactionID, - }) + result, err := h.OpenInteraction(ctx, agentexecution.OpenInteractionInput{ToolRef: agentexecution.ToolRef{ + ID: t.input.ToolCallID, + RuntimeLockID: t.input.RuntimeLockID}, + Question: &input.Form, + Destination: destination}) if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("load created question interaction: %w", err) - } - t.hasDurableCompletionOwner = true - record := agentInteractionRecordFromSQLC(row) - if err := handleSubagentQuestionTx(ctx, t.notifications, t.tx, t.q, record); err != nil { return AgentInteractionRecord{}, err } - return record, nil -} - -func markToolCallAwaitingPermissionTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - projectID, agentID, toolCallID, runtimeLockID uuid.UUID, - interactionUpdate *notifications.AgentInteractionUpdate, -) error { - _, err := qtx.MarkToolCallAwaitingPermission( + row, err := t.q.GetAgentInteraction( ctx, - dbsqlc.MarkToolCallAwaitingPermissionParams{ - ProjectID: projectID, - AgentID: agentID, - ID: toolCallID, - RuntimeLockID: runtimeLockID, + dbsqlc.GetAgentInteractionParams{ + ProjectID: t.input.ProjectID, + AgentID: t.input.AgentID, + ID: result.ID, }, ) - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return err - } - if err == nil { - txNotifications.AddToolCallUpdate( - agentID, - toolCallID, - string(ToolCallStateAwaitingPermission), - interactionUpdate, - ) - return nil + if err != nil { + return AgentInteractionRecord{}, err } - if runtimeErr := ensureRuntimeLockActiveTx( - ctx, - tx, - projectID, - agentID, - runtimeLockID, - ); runtimeErr != nil { - return runtimeErr + t.hasDurableCompletionOwner = true + t.disposition = ToolCallDispositionWaiting + t.applied = result.Created + if result.State != "open" { + t.disposition = ToolCallDispositionCompleted } - return storeerr.ErrIdempotencyConflict + return agentInteractionRecordFromSQLC(row), nil } func (s *Store) ResolveAgentInteraction( @@ -287,536 +155,56 @@ func (s *Store) ResolveAgentInteraction( if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ID == uuid.Nil { return AgentInteractionRecord{}, errors.New("project, agent, and interaction are required") } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentInteractionRecord{}, fmt.Errorf("begin resolve agent interaction: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - record, err := resolveAgentInteractionTx(ctx, txNotifications, tx, qtx, input) + defer func() { _ = unit.Rollback(ctx) }() + record, err := resolveExecutionInteraction(ctx, unit, input) if err != nil { return AgentInteractionRecord{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "resolve agent interaction"); err != nil { + if err := unit.Commit(ctx, "resolve agent interaction"); err != nil { return AgentInteractionRecord{}, err } return record, nil } -func resolveAgentInteractionTx( +func resolveExecutionInteraction( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, + unit *agentexecution.Unit, input ResolveAgentInteractionInput, ) (AgentInteractionRecord, error) { - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - ); err != nil { - return AgentInteractionRecord{}, fmt.Errorf("lock agent for interaction resolution: %w", err) - } - existing, err := qtx.GetAgentInteraction( - ctx, - dbsqlc.GetAgentInteractionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - }, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - return AgentInteractionRecord{}, fmt.Errorf("get agent interaction: %w", err) - } - resolution, err := normalizeAgentInteractionResolution( - AgentInteractionKind(existing.InteractionKind), - existing.Request, - input.Resolution, - ) - if err != nil { + if err := unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, + AgentID: input.AgentID}}, + agentexecution.ExternalAuthority{}); err != nil { return AgentInteractionRecord{}, err } - var record AgentInteractionRecord - if AgentInteractionState(existing.State) == AgentInteractionStateOpen { - stopped, err := qtx.AgentInteractionHasLaterStop( - ctx, - dbsqlc.AgentInteractionHasLaterStopParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: input.ID}, - ) - if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("check interaction stop lineage: %w", err) - } - if stopped { - return AgentInteractionRecord{}, storeerr.ErrStateTransitionConflict - } - resolvedByActorID, err := resolveActorTx( - ctx, - qtx, - input.ProjectID, - input.Actor, - ) - if err != nil { - return AgentInteractionRecord{}, err - } - responseInput, err := insertInteractionResponseInputTx(ctx, qtx, input, resolvedByActorID) - if err != nil { - return AgentInteractionRecord{}, err - } - eventRecord, err := appendTypedAgentEventTx(ctx, txNotifications, tx, AppendTypedAgentEventInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: existing.TurnID, - Kind: events.KindAgentInput, - IdempotencyKey: "agent_input:" + responseInput.ID.String(), - AgentInputID: responseInput.ID, - }) - if err != nil { - return AgentInteractionRecord{}, err - } - if err := updateAgentTurnLatestEventTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - existing.TurnID, - eventRecord.Event.ID, - eventRecord.Event.ID, - ); err != nil { - return AgentInteractionRecord{}, err - } - _, err = qtx.ResolveAgentInteraction( - ctx, - dbsqlc.ResolveAgentInteractionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - Resolution: resolution, - ResolvedByInputID: storeutil.IDFromNil(responseInput.ID), - }, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - return AgentInteractionRecord{}, fmt.Errorf("resolve agent interaction: %w", err) - } - row, err := qtx.GetAgentInteraction(ctx, dbsqlc.GetAgentInteractionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - }) - if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("load resolved agent interaction: %w", err) - } - if _, err := qtx.ResolveInteractionResponseAgentInput( - ctx, - dbsqlc.ResolveInteractionResponseAgentInputParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: responseInput.ID, - TargetInteractionID: storeutil.IDFromNil(input.ID), - EventID: storeutil.IDFromNil(eventRecord.Event.ID), - }, - ); err != nil { - return AgentInteractionRecord{}, fmt.Errorf("resolve interaction response agent input: %w", err) - } - record = agentInteractionRecordFromSQLC(row) - if err := applyPermissionInteractionResolutionTx(ctx, txNotifications, tx, qtx, record); err != nil { - return AgentInteractionRecord{}, err - } - if err := completeQuestionToolCallTx(ctx, txNotifications, tx, qtx, record); err != nil { - return AgentInteractionRecord{}, err - } - } else { - record = agentInteractionRecordFromSQLC(existing) - if record.State != AgentInteractionStateResolved || !sameJSON(record.Resolution, resolution) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - requestActorID, actorFound, err := lookupActorIDTx(ctx, qtx, input.ProjectID, input.Actor) - if err != nil { - return AgentInteractionRecord{}, err - } - responseInput, responseFound, err := loadAgentInputByIdempotencyMaybeTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - interactionResponseIdempotencyScope, - input.ID.String(), - ) - if err != nil { - return AgentInteractionRecord{}, err - } - if !actorFound || !responseFound || - record.ResolvedByInputID != responseInput.ID || - responseInput.ActorID != requestActorID { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - if err := completeQuestionToolCallTx(ctx, txNotifications, tx, qtx, record); err != nil { - return AgentInteractionRecord{}, err - } - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - Metadata: []byte(`{"reason":"interaction_resolved"}`), - }, - ); err != nil { - return AgentInteractionRecord{}, fmt.Errorf("mark interaction resolution wakeup: %w", err) - } - return record, nil -} - -func applyPermissionInteractionResolutionTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - record AgentInteractionRecord, -) error { - if record.InteractionKind != AgentInteractionKindPermission { - return nil - } - resolution, err := permissionInteractionResolution(record) + q := dbsqlc.New(unit.DB()) + actor, err := resolveActorTx(ctx, q, input.ProjectID, input.Actor) if err != nil { - return err - } - if resolution.Decision == toolpermission.DecisionAllow { - row, err := qtx.MarkToolCallReadyFromInteraction( - ctx, - dbsqlc.MarkToolCallReadyFromInteractionParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - InteractionID: record.ID, - }, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - existing, loadErr := qtx.GetToolCall( - ctx, - dbsqlc.GetToolCallParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - ID: record.ToolCallID, - }, - ) - if loadErr != nil { - return loadErr - } - if existing.State == string(ToolCallStateReady) || - existing.State == string(ToolCallStateRunning) || - existing.State == string(ToolCallStateWaiting) || - (existing.State == string(ToolCallStateCompleted) && - existing.Outcome != string(ToolResultOutcomeDenied) && - existing.Outcome != string(ToolResultOutcomeCanceled)) { - return nil - } - return storeerr.ErrIdempotencyConflict - } - return fmt.Errorf("mark permitted tool call ready: %w", err) - } - txNotifications.AddToolCallUpdate(row.AgentID, row.ID, row.State, interactionUpdateFromRecord(record)) - return nil - } - return completePermissionDeniedToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - resolution.Reason, - ) -} - -func completePermissionDeniedToolCallTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - record AgentInteractionRecord, - reason string, -) error { - result, err := marshalJSON(map[string]any{"reason": reason}) - if err != nil { - return fmt.Errorf("marshal permission denied tool result: %w", err) + return AgentInteractionRecord{}, err } - contentParts, err := ToolResultContentParts(result) + h, err := unit.Handle(input.ProjectID, input.AgentID) if err != nil { - return err + return AgentInteractionRecord{}, err } - outcome := permissionToolResultOutcome(record) - row, err := qtx.CompleteToolCallFromPermissionInteraction( + result, err := h.ResolveInteraction( ctx, - dbsqlc.CompleteToolCallFromPermissionInteractionParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - InteractionID: record.ID, - Outcome: string(outcome), - }, + agentexecution.ResolveInteractionInput{ID: input.ID, ActorID: actor, Resolution: input.Resolution}, ) - if errors.Is(err, pgx.ErrNoRows) { - existing, loadErr := qtx.GetToolCall( - ctx, - dbsqlc.GetToolCallParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - ID: record.ToolCallID, - }, - ) - if loadErr != nil { - return loadErr - } - if existing.State != string(ToolCallStateCompleted) || - existing.Outcome != string(outcome) { - return storeerr.ErrIdempotencyConflict - } - ok, checkErr := completedToolCallMatchesTx( - ctx, - qtx, - record.ProjectID, - record.AgentID, - record.ToolCallID, - outcome, - contentParts, - ) - if checkErr != nil { - return checkErr - } - if ok { - return nil - } - return storeerr.ErrIdempotencyConflict - } - if err != nil { - return fmt.Errorf("complete tool call from permission interaction: %w", err) - } - resultRecord := toolCallRecordFromPermissionInteractionCompleteSQLC(row) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx( - ctx, txNotifications, tx, resultRecord, interactionUpdateFromRecord(record), - ); err != nil { - return err - } - return nil -} - -func completeQuestionToolCallTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - interaction AgentInteractionRecord, -) error { - if interaction.InteractionKind != AgentInteractionKindQuestion { - return nil - } - outcome := ToolResultOutcomeSucceeded - var result json.RawMessage - var err error - switch interaction.State { - case AgentInteractionStateResolved: - form, parseErr := interaction.Form() - if parseErr != nil { - return parseErr - } - response, parseErr := interactionform.ParseResolution(form, interaction.Resolution) - if parseErr != nil { - return parseErr - } - result, err = marshalJSON(newQuestionToolResult(form, response)) - case AgentInteractionStateCanceled: - outcome = ToolResultOutcomeCanceled - result, err = marshalJSON(map[string]any{"reason": "question interaction canceled"}) - case AgentInteractionStateOpen: - return nil - default: - return fmt.Errorf( - "question interaction %s has unsupported state %q", - interaction.ID, - interaction.State, - ) - } - if err != nil { - return fmt.Errorf("marshal question tool result: %w", err) - } - contentParts, err := ToolResultContentParts(result) if err != nil { - return err - } - if len(contentParts) > ToolResultInlineBudgetBytes { - ok, err := completedToolCallMatchesTx( - ctx, qtx, interaction.ProjectID, interaction.AgentID, interaction.ToolCallID, outcome, contentParts, - ) - if err != nil { - return err - } - if ok { - return nil - } - outcome = ToolResultOutcomeFailed - contentParts, err = ToolResultContentParts(json.RawMessage( - `{"error":"Question response exceeded the size limit. Ask fewer or shorter questions, or request a shorter answer."}`, - )) - if err != nil { - return err - } - } - row, err := qtx.CompleteToolCallFromQuestionInteraction( - ctx, - dbsqlc.CompleteToolCallFromQuestionInteractionParams{ - ProjectID: interaction.ProjectID, - AgentID: interaction.AgentID, - InteractionID: interaction.ID, - Outcome: string(outcome), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - ok, checkErr := completedToolCallMatchesTx( - ctx, - qtx, - interaction.ProjectID, - interaction.AgentID, - interaction.ToolCallID, - outcome, - contentParts, - ) - if checkErr != nil { - return checkErr - } - if ok { - return nil - } - return storeerr.ErrIdempotencyConflict - } - if err != nil { - return fmt.Errorf("complete tool call from question interaction: %w", err) - } - resultRecord := toolCallRecordFromQuestionInteractionCompleteSQLC(row) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx( - ctx, txNotifications, tx, resultRecord, interactionUpdateFromRecord(interaction), - ); err != nil { - return err - } - return nil -} - -type questionToolResult struct { - Answers []questionToolResultAnswer `json:"answers"` -} - -type questionToolResultAnswer struct { - QuestionIndex int `json:"question_index"` - Question string `json:"question"` - SelectedOptions []questionToolResultSelectedOption `json:"selected_options"` - Text string `json:"text,omitempty"` -} - -type questionToolResultSelectedOption struct { - OptionIndex int `json:"option_index"` - Label string `json:"label"` -} - -func newQuestionToolResult( - form interactionform.Form, - resolution interactionform.Resolution, -) questionToolResult { - result := questionToolResult{ - Answers: make([]questionToolResultAnswer, 0, len(resolution.Answers)), - } - for questionIndex, answer := range resolution.Answers { - question := form.Questions[questionIndex] - selectedOptions := make( - []questionToolResultSelectedOption, - 0, - len(answer.OptionIndices), - ) - for _, optionIndex := range answer.OptionIndices { - selectedOptions = append(selectedOptions, questionToolResultSelectedOption{ - OptionIndex: optionIndex, - Label: question.Options[optionIndex].Label, - }) - } - result.Answers = append(result.Answers, questionToolResultAnswer{ - QuestionIndex: questionIndex, - Question: question.Prompt, - SelectedOptions: selectedOptions, - Text: answer.Text, - }) + return AgentInteractionRecord{}, err } - return result -} - -func insertInteractionResponseInputTx( - ctx context.Context, - qtx *dbsqlc.Queries, - input ResolveAgentInteractionInput, - resolvedByActorID uuid.UUID, -) (AgentInputRecord, error) { - scope := interactionResponseIdempotencyScope - key := input.ID.String() - row, err := qtx.InsertInteractionResponseAgentInput( + row, err := q.GetAgentInteraction( ctx, - dbsqlc.InsertInteractionResponseAgentInputParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TargetInteractionID: input.ID, - ActorID: storeutil.IDFromNil(resolvedByActorID), - IdempotencyScope: storeutil.TextFromEmpty(scope), - InputIdempotencyKey: storeutil.TextFromEmpty(key), - Metadata: json.RawMessage(`{}`), - }, + dbsqlc.GetAgentInteractionParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: result.ID}, ) - if errors.Is(err, pgx.ErrNoRows) { - existingInput, getErr := qtx.GetAgentInputByIdempotency( - ctx, - dbsqlc.GetAgentInputByIdempotencyParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - IdempotencyScope: scope, - InputIdempotencyKey: key, - }, - ) - if getErr != nil { - return AgentInputRecord{}, fmt.Errorf( - "load interaction response input by idempotency: %w", - getErr, - ) - } - record := agentInputRecordFromIdempotencySQLC(existingInput) - if record.InputKind != "interaction_response" || record.TargetInteractionID != input.ID { - return AgentInputRecord{}, storeerr.ErrIdempotencyConflict - } - return record, nil - } - if err != nil { - return AgentInputRecord{}, fmt.Errorf("insert interaction response agent input: %w", err) - } - return agentInputRecordFromInteractionResponseInsertSQLC(row), nil -} - -func normalizeAgentInteractionResolution( - kind AgentInteractionKind, - request json.RawMessage, - resolution interactionform.Resolution, -) (json.RawMessage, error) { - value, err := interactionFormForInteraction(kind, request) if err != nil { - return nil, err - } - normalized, err := interactionform.NormalizeResolution(value, resolution) - if err != nil { - return nil, err - } - raw, err := marshalJSON(normalized) - if err != nil { - return nil, fmt.Errorf("marshal interaction resolution: %w", err) + return AgentInteractionRecord{}, err } - return raw, nil + return agentInteractionRecordFromSQLC(row), nil } func interactionFormForInteraction( @@ -850,42 +238,6 @@ func interactionFormForInteraction( } } -func permissionInteractionResolution( - record AgentInteractionRecord, -) (toolpermission.Resolution, error) { - if record.InteractionKind != AgentInteractionKindPermission { - return toolpermission.Resolution{}, errors.New("permission interaction is required") - } - if record.State == AgentInteractionStateCanceled { - return toolpermission.Resolution{ - Decision: toolpermission.DecisionDeny, - Reason: "permission interaction canceled", - }, nil - } - if record.State != AgentInteractionStateResolved { - return toolpermission.Resolution{}, fmt.Errorf( - "permission interaction has unsupported state %q", - record.State, - ) - } - request, err := toolpermission.ParseRequest(record.Request) - if err != nil { - return toolpermission.Resolution{}, err - } - resolution, err := interactionform.ParseResolution(request.Form, record.Resolution) - if err != nil { - return toolpermission.Resolution{}, err - } - return toolpermission.Resolve(request, resolution) -} - -func permissionToolResultOutcome(record AgentInteractionRecord) ToolResultOutcome { - if record.State == AgentInteractionStateCanceled { - return ToolResultOutcomeCanceled - } - return ToolResultOutcomeDenied -} - type ListAgentInteractionsForAgentInput struct { ProjectID uuid.UUID AgentID uuid.UUID diff --git a/internal/storage/executionstore/agent_interactions_store_test.go b/internal/storage/executionstore/agent_interactions_store_test.go deleted file mode 100644 index 54bbafb5a..000000000 --- a/internal/storage/executionstore/agent_interactions_store_test.go +++ /dev/null @@ -1,125 +0,0 @@ -package executionstore - -import ( - "encoding/json" - "testing" - - "github.com/omnara-ai/omnara/internal/interactionform" - "github.com/omnara-ai/omnara/internal/toolpermission" -) - -func TestPermissionInteractionResolutionDelegatesToPermissionMode(t *testing.T) { - request := permissionRequestForStorageTest(t, "test_tool") - for _, test := range []struct { - name string - state AgentInteractionState - resolution interactionform.Resolution - wantDecision toolpermission.Decision - }{ - { - name: "allowed", - state: AgentInteractionStateResolved, - resolution: interactionform.Resolution{Answers: []interactionform.Answer{{ - OptionIndices: []int{toolpermission.AllowOptionIndex}, - }}}, - wantDecision: toolpermission.DecisionAllow, - }, - { - name: "denied", - state: AgentInteractionStateResolved, - resolution: interactionform.Resolution{Answers: []interactionform.Answer{{ - OptionIndices: []int{toolpermission.DenyOptionIndex}, - Text: "not allowed", - }}}, - wantDecision: toolpermission.DecisionDeny, - }, - { - name: "canceled", - state: AgentInteractionStateCanceled, - wantDecision: toolpermission.DecisionDeny, - }, - } { - t.Run(test.name, func(t *testing.T) { - resolution, err := json.Marshal(test.resolution) - if err != nil { - t.Fatalf("marshal resolution: %v", err) - } - record := AgentInteractionRecord{ - InteractionKind: "permission", - State: test.state, - Request: request, - Resolution: resolution, - } - decision, err := permissionInteractionResolution(record) - if err != nil { - t.Fatalf("permissionInteractionResolution() error = %v", err) - } - if decision.Decision != test.wantDecision { - t.Fatalf("decision = %q, want %q", decision.Decision, test.wantDecision) - } - }) - } -} - -func TestQuestionToolResultKeepsPositionalIdentityForDuplicateLabels(t *testing.T) { - result := newQuestionToolResult( - interactionform.Form{Questions: []interactionform.Question{{ - Prompt: "Choose", - Multiple: true, - Options: []interactionform.Option{ - {Label: "Other"}, - {Label: "Other", AllowsText: true}, - }, - }}}, - interactionform.Resolution{Answers: []interactionform.Answer{{ - OptionIndices: []int{0, 1}, - }}}, - ) - if len(result.Answers) != 1 || result.Answers[0].QuestionIndex != 0 { - t.Fatalf("question result = %+v", result) - } - selected := result.Answers[0].SelectedOptions - if len(selected) != 2 || - selected[0].OptionIndex != 0 || - selected[1].OptionIndex != 1 || - selected[0].Label != "Other" || - selected[1].Label != "Other" { - t.Fatalf("selected options = %+v", selected) - } -} - -func permissionRequestForStorageTest(t *testing.T, toolName string) json.RawMessage { - t.Helper() - authorization, err := toolpermission.NewAuthorization( - toolName, - json.RawMessage(`{}`), - ) - if err != nil { - t.Fatalf("authorization: %v", err) - } - mode, ok := toolpermission.FindMode( - toolpermission.CommonModeDescriptors(), - toolpermission.ModeAlwaysAsk, - ) - if !ok { - t.Fatal("always_ask mode missing") - } - value, err := toolpermission.NewAllowDenyForm("Permission requested", nil) - if err != nil { - t.Fatalf("interaction form: %v", err) - } - request, err := toolpermission.NewRequest( - mode, - toolpermission.DefaultSelection(toolpermission.ModeAlwaysAsk), - authorization, - value, - ) - if err != nil { - t.Fatalf("request: %v", err) - } - raw, err := json.Marshal(request) - if err != nil { - t.Fatalf("marshal request: %v", err) - } - return raw -} diff --git a/internal/storage/executionstore/agent_launch_inbox.go b/internal/storage/executionstore/agent_launch_inbox.go index 0e64e78e9..7b5540396 100644 --- a/internal/storage/executionstore/agent_launch_inbox.go +++ b/internal/storage/executionstore/agent_launch_inbox.go @@ -6,10 +6,10 @@ import ( "errors" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -84,12 +84,14 @@ func (s *Store) admitInboxLaunchRecipientOnce( ); err != nil || outcome.Outcome != InboxRecipientPending { return LaunchAgentResult{Agent: outcome.Agent}, err } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return LaunchAgentResult{}, err } - defer func() { _ = tx.Rollback(ctx) }() - q := s.q.WithTx(tx) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + q := dbsqlc.New(tx) resources, err := launchIntegrationIDsTx(ctx, q, recipient.Launch.launchInput(lease.ProjectID, recipient.InitialInput)) if err != nil { return LaunchAgentResult{}, err @@ -97,7 +99,7 @@ func (s *Store) admitInboxLaunchRecipientOnce( work, err := s.integrations.LockIntegrationInboxLeaseTx(ctx, tx, lease, resources...) if err != nil { // Release the connection before diagnostic reads, which may need the pool's only session. - _ = tx.Rollback(ctx) + _ = unit.Rollback(ctx) if outcome, readErr := resolveInboxLaunchOutcome( ctx, s.q, snapshot.ProjectID, recipient, ); readErr == nil && outcome.Outcome != InboxRecipientPending { @@ -148,8 +150,7 @@ func (s *Store) admitInboxLaunchRecipientOnce( LaunchKey: claim.LaunchKey, Artifacts: artifacts, } - txNotifications := s.newTxNotifications() - result, err := s.launchAgentTx(ctx, tx, q, txNotifications, launch) + result, err := s.launchAgentTx(ctx, unit, q, launch) if err != nil { return LaunchAgentResult{}, err } @@ -159,7 +160,7 @@ func (s *Store) admitInboxLaunchRecipientOnce( if err := work.CheckLease(ctx); err != nil { return LaunchAgentResult{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "admit inbox launch recipient"); err != nil { + if err := unit.Commit(ctx, "admit inbox launch recipient"); err != nil { return LaunchAgentResult{}, err } return result, nil diff --git a/internal/storage/executionstore/agent_launch_input.go b/internal/storage/executionstore/agent_launch_input.go index 5c710c1c8..90b0f53ac 100644 --- a/internal/storage/executionstore/agent_launch_input.go +++ b/internal/storage/executionstore/agent_launch_input.go @@ -7,10 +7,9 @@ import ( "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -83,8 +82,7 @@ func prepareLaunchInitialInput(input LaunchAgentInput) (*LaunchInitialInput, []C func (s *Store) insertLaunchInitialContentInputTx( ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, + unit *agentexecution.Unit, agent AgentRecord, launch LaunchAgentInput, initial LaunchInitialInput, @@ -92,6 +90,7 @@ func (s *Store) insertLaunchInitialContentInputTx( admission *launchAdmission, result *LaunchAgentResult, ) error { + tx := unit.DB() q := dbsqlc.New(tx) actor := initial.Actor var err error @@ -153,7 +152,7 @@ func (s *Store) insertLaunchInitialContentInputTx( return err } } - created, err := createAgentContentInputTx(ctx, txNotifications, tx, q, content, blocks) + created, err := createAgentContentInputTx(ctx, unit, q, content, blocks) if err != nil { return err } diff --git a/internal/storage/executionstore/agent_launch_integration_test.go b/internal/storage/executionstore/agent_launch_integration_test.go index 79dea3c7a..2799bf013 100644 --- a/internal/storage/executionstore/agent_launch_integration_test.go +++ b/internal/storage/executionstore/agent_launch_integration_test.go @@ -270,7 +270,10 @@ tools: } if len(result.MachineBindings) != 1 || len(result.ProvisionMachineIDs) != 1 || result.ProvisionMachineIDs[0] != result.MachineBindings[0].MachineID { - t.Fatalf("pool launch should create one attached binding and provisioning request, result=%+v", result) + t.Fatalf( + "pool launch should create one attached binding and provisioning request, result=%+v", + result, + ) } return result } @@ -379,7 +382,10 @@ func TestConcurrentSameKeyLaunchIgnoresLosingBody(t *testing.T) { } } -func testConcurrentSameKeyLaunchIgnoresLosingBody(t *testing.T, mutate func(*executionstore.LaunchAgentInput)) { +func testConcurrentSameKeyLaunchIgnoresLosingBody( + t *testing.T, + mutate func(*executionstore.LaunchAgentInput), +) { t.Helper() ctx := context.Background() pool := openIntegrationDB(t, ctx) @@ -577,7 +583,13 @@ func TestLaunchAgentWithDefaultPool(t *testing.T) { DefaultMachineProviderOptions: json.RawMessage(`{"image":"default"}`), }, )) - user := mustCreateProjectDeveloperUser(t, ctx, store, "launch-default-pool@example.com", "Launch Cluster Pool") + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "launch-default-pool@example.com", + "Launch Cluster Pool", + ) if err := provisionDefaultMachinePoolGrantsForProject( ctx, store, @@ -668,7 +680,9 @@ tools: t.Fatalf("default pool project machine pool grants = %d, want 1", poolGrantCount) } var defaultPoolGrantID uuid.UUID - if err := pool.QueryRow(ctx, `SELECT id FROM project_machine_pool_grants WHERE project_id = $1`, testProjectID). + if err := pool.QueryRow(ctx, + `SELECT id FROM project_machine_pool_grants WHERE project_id = $1`, + testProjectID). Scan(&defaultPoolGrantID); err != nil { t.Fatalf("load default project machine pool grant: %v", err) } @@ -682,7 +696,8 @@ tools: testProjectID, binding.MachineID, ) - if generatedGrant.SourceKind != "pool" || generatedGrant.ProjectMachinePoolGrantID != defaultPoolGrantID { + if generatedGrant.SourceKind != "pool" || + generatedGrant.ProjectMachinePoolGrantID != defaultPoolGrantID { t.Fatalf( "generated grant source_kind=%q project_machine_pool_grant_id=%s, want pool %s", generatedGrant.SourceKind, @@ -808,7 +823,11 @@ func TestArchiveAgentMarksPoolMachinesDeletingAndStopsExecution(t *testing.T) { } backlogBeforeArchive, err := store.Execution().ListQueuedBacklogInputs( ctx, - executionstore.ListQueuedBacklogInputsInput{ProjectID: testProjectID, AgentID: result.Agent.ID, Limit: 10}, + executionstore.ListQueuedBacklogInputsInput{ + ProjectID: testProjectID, + AgentID: result.Agent.ID, + Limit: 10, + }, ) if err != nil { t.Fatalf("list queued backlog before archive: %v", err) @@ -959,7 +978,11 @@ WHERE agent.project_id = $1 } backlogAfterArchive, err := store.Execution().ListQueuedBacklogInputs( ctx, - executionstore.ListQueuedBacklogInputsInput{ProjectID: testProjectID, AgentID: result.Agent.ID, Limit: 10}, + executionstore.ListQueuedBacklogInputsInput{ + ProjectID: testProjectID, + AgentID: result.Agent.ID, + Limit: 10, + }, ) if err != nil { t.Fatalf("list queued backlog after archive: %v", err) @@ -984,11 +1007,10 @@ WHERE agent.project_id = $1 archivedAt, ) } - if err := store.Execution().MarkAgentWakeup( + if err := store.Execution().RepairExecutionForTest( ctx, testProjectID, result.Agent.ID, - json.RawMessage(`{"reason":"after_archive"}`), ); err != nil { t.Fatalf("mark archived wakeup: %v", err) } @@ -1023,7 +1045,13 @@ func TestChangeAgentConfigReconcilesPoolMachineSources(t *testing.T) { seedMigratedDB(t, ctx, pool) store := newIntegrationStore(pool, storage.WithMachinePoolProviders(mergingMachinePoolProviders{})) now := time.Date(2026, 5, 21, 9, 48, 0, 0, time.UTC) - user := mustCreateProjectDeveloperUser(t, ctx, store, "live-pool-sources@example.com", "Live Pool Sources") + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "live-pool-sources@example.com", + "Live Pool Sources", + ) machinePool := createLaunchTestMachinePool( t, ctx, @@ -1038,12 +1066,14 @@ func TestChangeAgentConfigReconcilesPoolMachineSources(t *testing.T) { 5, now, ) - if _, err := store.Execution().CreateProjectMachinePoolGrant(ctx, executionstore.CreateProjectMachinePoolGrantInput{ - OrgID: testOrgID, - ProjectID: testProjectID, - MachinePoolID: machinePool.ID, - IdempotencyKey: "idem-live-pool-grant", - }); err != nil { + if _, + err := store.Execution().CreateProjectMachinePoolGrant(ctx, + executionstore.CreateProjectMachinePoolGrantInput{ + OrgID: testOrgID, + ProjectID: testProjectID, + MachinePoolID: machinePool.ID, + IdempotencyKey: "idem-live-pool-grant", + }); err != nil { t.Fatalf("create live pool grant: %v", err) } emptyYAML := ` @@ -1099,7 +1129,12 @@ tools: if len(added.DeleteMachines) != 0 { t.Fatalf("added pool source deleted machines: %+v", added.DeleteMachines) } - machines, err := executionstore.IntegrationListPoolMachinesTx(ctx, store.q, testProjectID, launch.Agent.ID) + machines, err := executionstore.IntegrationListPoolMachinesTx( + ctx, + store.q, + testProjectID, + launch.Agent.ID, + ) if err != nil { t.Fatalf("list added pool machines: %v", err) } @@ -1374,7 +1409,8 @@ tools: if err != nil { t.Fatalf("load removed pool machine: %v", err) } - if deleted.LifecycleState != "deleting" || deleted.LifecycleReasonCode != "agent_config_machine_source_removed" { + if deleted.LifecycleState != "deleting" || + deleted.LifecycleReasonCode != "agent_config_machine_source_removed" { t.Fatalf("removed pool machine = %+v", deleted) } assertProviderRuntimeMismatchClearedForTest( @@ -1396,7 +1432,11 @@ tools: } if replayedRemoval.ConfigChange.AgentInput.ID != removed.ConfigChange.AgentInput.ID || replayedRemoval.ConfigChange.Event.ID != removed.ConfigChange.Event.ID { - t.Fatalf("removed pool config replay = %+v, want %+v", replayedRemoval.ConfigChange, removed.ConfigChange) + t.Fatalf( + "removed pool config replay = %+v, want %+v", + replayedRemoval.ConfigChange, + removed.ConfigChange, + ) } if len(replayedRemoval.DeleteMachines) != len(removed.DeleteMachines) { t.Fatalf( @@ -1503,7 +1543,8 @@ tools: json.RawMessage(compiled.CanonicalJSON), compiled.Hash) - if err == nil || !strings.Contains(err.Error(), "secret_env.API_TOKEN secret is not available to the project") { + if err == nil || + !strings.Contains(err.Error(), "secret_env.API_TOKEN secret is not available to the project") { t.Fatalf("default pool secret_env validation error = %v", err) } if _, err := store.Secrets().CreateSecretGrant(ctx, secretstore.CreateSecretGrantInput{ @@ -1817,7 +1858,10 @@ func TestLaunchAgentCreatesMultipleMachineBindings(t *testing.T) { store := newIntegrationStore(pool, storage.WithMachinePoolProviders(mergingMachinePoolProviders{})) user, err := store.Identity().CreateVerifiedUser( ctx, - storagetest.CreateVerifiedUserInput{Email: "launch-multi@example.com", DisplayName: "Launch Multi User"}, + storagetest.CreateVerifiedUserInput{ + Email: "launch-multi@example.com", + DisplayName: "Launch Multi User", + }, ) if err != nil { t.Fatalf("create user: %v", err) @@ -1868,7 +1912,13 @@ func TestLaunchAgentCreatesMultipleMachineBindings(t *testing.T) { if err != nil { t.Fatalf("create second grant: %v", err) } - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "launch-multi", "Multi Launch Agent", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "launch-multi", + "Multi Launch Agent", + ` instruction: Use the selected machines. model: provider_config: openai-prod @@ -1882,7 +1932,8 @@ machine_sources: description: Second machine tools: run_command: {} -`) +`, + ) result, err := store.Execution().LaunchAgent( ctx, @@ -1954,7 +2005,8 @@ tools: t.Fatalf("register daemon runtime %d: %v", index, err) } } - executable, err := store.Execution().ListExecutableAgentMachineBindings(ctx, testProjectID, result.Agent.ID) + executable, err := store.Execution(). + ListExecutableAgentMachineBindings(ctx, testProjectID, result.Agent.ID) if err != nil { t.Fatalf("list executable bindings: %v", err) } @@ -2116,7 +2168,12 @@ tools: t.Fatalf("unexpected generated grant %d: %+v binding=%+v", slotIndex, generatedGrant, binding) } } - poolMachines, err := executionstore.IntegrationListPoolMachinesTx(ctx, store.q, testProjectID, result.Agent.ID) + poolMachines, err := executionstore.IntegrationListPoolMachinesTx( + ctx, + store.q, + testProjectID, + result.Agent.ID, + ) if err != nil { t.Fatalf("list launched pool machines: %v", err) } @@ -2227,7 +2284,8 @@ tools: if len(firstLaunch.MachineBindings) != 1 { t.Fatalf("first launch bindings = %+v, want one", firstLaunch.MachineBindings) } - firstMachine, err := store.Execution().GetMachine(ctx, testOrgID, firstLaunch.MachineBindings[0].MachineID) + firstMachine, err := store.Execution(). + GetMachine(ctx, testOrgID, firstLaunch.MachineBindings[0].MachineID) if err != nil { t.Fatalf("get first generated machine: %v", err) } @@ -2305,7 +2363,8 @@ tools: secondGrant.ID, ) } - secondMachine, err := store.Execution().GetMachine(ctx, testOrgID, secondLaunch.MachineBindings[0].MachineID) + secondMachine, err := store.Execution(). + GetMachine(ctx, testOrgID, secondLaunch.MachineBindings[0].MachineID) if err != nil { t.Fatalf("get second generated machine: %v", err) } @@ -2410,7 +2469,13 @@ func TestLaunchAgentMultiplePoolSourcesKeepSourceOrder(t *testing.T) { if err != nil { t.Fatalf("create second pool grant: %v", err) } - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "launch-two-pools", "Launch Two Pools Agent", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "launch-two-pools", + "Launch Two Pools Agent", + ` instruction: Use independent pool machines. model: provider_config: openai-prod @@ -2428,7 +2493,8 @@ machine_sources: description: Second pool machine tools: run_command: {} -`) +`, + ) result, err := store.Execution().LaunchAgent( ctx, @@ -2444,7 +2510,11 @@ tools: t.Fatalf("launch two-pool agent: %v", err) } if len(result.MachineBindings) != 4 || len(result.ProvisionMachineIDs) != 4 { - t.Fatalf("two-pool launch result bindings=%+v provision=%v", result.MachineBindings, result.ProvisionMachineIDs) + t.Fatalf( + "two-pool launch result bindings=%+v provision=%v", + result.MachineBindings, + result.ProvisionMachineIDs, + ) } want := []struct { cwd string @@ -2546,7 +2616,13 @@ func TestLaunchAgentZeroInitialPoolValidatesGrantWithoutCreatingMachines(t *test }); err != nil { t.Fatalf("create pool grant: %v", err) } - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "launch-zero-pool", "Launch Zero Pool Agent", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "launch-zero-pool", + "Launch Zero Pool Agent", + ` instruction: Use a pool later. model: provider_config: openai-prod @@ -2556,7 +2632,8 @@ machine_sources: max_machines: 0 tools: run_command: {} -`) +`, + ) result, err := store.Execution().LaunchAgent( ctx, @@ -2578,7 +2655,8 @@ tools: result.ProvisionMachineIDs, ) } - executable, err := store.Execution().ListExecutableAgentMachineBindings(ctx, testProjectID, result.Agent.ID) + executable, err := store.Execution(). + ListExecutableAgentMachineBindings(ctx, testProjectID, result.Agent.ID) if err != nil { t.Fatalf("list zero-initial executable bindings: %v", err) } @@ -2701,7 +2779,9 @@ func TestLaunchAgentZeroInitialPoolSkipsCapacityCheck(t *testing.T) { store, "Launch Zero Capacity Pool", "test", - defaultMachineFieldsForTest{DefaultMachineProviderOptions: json.RawMessage(`{"image":"zero-capacity"}`)}, + defaultMachineFieldsForTest{ + DefaultMachineProviderOptions: json.RawMessage(`{"image":"zero-capacity"}`), + }, 1, now, ) @@ -2795,7 +2875,9 @@ func TestLaunchAgentInitialPoolCapacityRollsBackAllRows(t *testing.T) { store, "Launch Initial Capacity Pool", "test", - defaultMachineFieldsForTest{DefaultMachineProviderOptions: json.RawMessage(`{"image":"initial-capacity"}`)}, + defaultMachineFieldsForTest{ + DefaultMachineProviderOptions: json.RawMessage(`{"image":"initial-capacity"}`), + }, 1, now, ) @@ -3146,12 +3228,16 @@ tools: ); err != nil { t.Fatalf("revoke first pool grant: %v", err) } - firstMachine, err := store.Execution().GetMachine(ctx, testOrgID, firstLaunch.MachineBindings[0].MachineID) + firstMachine, err := store.Execution(). + GetMachine(ctx, testOrgID, firstLaunch.MachineBindings[0].MachineID) if err != nil { t.Fatalf("get first machine after revoke: %v", err) } if firstMachine.LifecycleState != "deleting" { - t.Fatalf("first machine lifecycle after grant revoke = %q, want deleting", firstMachine.LifecycleState) + t.Fatalf( + "first machine lifecycle after grant revoke = %q, want deleting", + firstMachine.LifecycleState, + ) } secondGrant, err := store.Execution().CreateProjectMachinePoolGrant( ctx, @@ -3381,7 +3467,9 @@ func TestConcurrentLaunchesRespectPoolCapacity(t *testing.T) { store, "Launch Concurrent Capacity Pool", "test", - defaultMachineFieldsForTest{DefaultMachineProviderOptions: json.RawMessage(`{"image":"concurrent-capacity"}`)}, + defaultMachineFieldsForTest{ + DefaultMachineProviderOptions: json.RawMessage(`{"image":"concurrent-capacity"}`), + }, 1, now, ) @@ -3699,7 +3787,14 @@ func TestPoolLaunchMachineProvisioningActivatesBindingAfterDaemonRuntime(t *test ); err == nil { t.Fatal("expected deletion completion before delete intent to fail") } - binding = getAgentMachineBindingForTest(t, ctx, store, testProjectID, result.Agent.ID, result.MachineBindings[0].ID) + binding = getAgentMachineBindingForTest( + t, + ctx, + store, + testProjectID, + result.Agent.ID, + result.MachineBindings[0].ID, + ) if binding.State != "attached" { t.Fatalf("rejected deletion completion changed binding: %+v", binding) } @@ -3788,7 +3883,14 @@ func TestPoolLaunchMachineProvisioningActivatesBindingAfterDaemonRuntime(t *test if deletedMachine.LifecycleState != "deleted" || deletedMachine.DeletedAt == nil { t.Fatalf("deleted machine = %+v", deletedMachine) } - binding = getAgentMachineBindingForTest(t, ctx, store, testProjectID, result.Agent.ID, result.MachineBindings[0].ID) + binding = getAgentMachineBindingForTest( + t, + ctx, + store, + testProjectID, + result.Agent.ID, + result.MachineBindings[0].ID, + ) if binding.State != "released" { t.Fatalf("binding after cleanup = %+v, want released", binding) } @@ -3856,7 +3958,11 @@ func TestPoolProvisionMaxAttemptsCleanupOnlyClaimsStaleProvisioning(t *testing.T next_reconcile_after = $4, updated_at = $4 WHERE org_id = $1 AND id = $2 - `, testOrgID, result.MachineBindings[0].MachineID, executionstore.DefaultPoolMachineProvisionFailureLimit, now.Add(2*time.Second)); err != nil { + `, + testOrgID, + result.MachineBindings[0].MachineID, + executionstore.DefaultPoolMachineProvisionFailureLimit, + now.Add(2*time.Second)); err != nil { t.Fatalf("seed exhausted provision failure: %v", err) } provisioning, err := store.Execution().ListPoolMachinesForProvisioning(ctx, 10) @@ -3884,7 +3990,11 @@ func TestPoolProvisionMaxAttemptsCleanupOnlyClaimsStaleProvisioning(t *testing.T next_reconcile_after = $4, updated_at = $4 WHERE org_id = $1 AND id = $2 - `, testOrgID, result.MachineBindings[0].MachineID, executionstore.DefaultPoolMachineProvisionFailureLimit, now.Add(-3*time.Minute)); err != nil { + `, + testOrgID, + result.MachineBindings[0].MachineID, + executionstore.DefaultPoolMachineProvisionFailureLimit, + now.Add(-3*time.Minute)); err != nil { t.Fatalf("seed exhausted stale provisioning: %v", err) } provisioning, err = store.Execution().ListPoolMachinesForProvisioning(ctx, 10) @@ -3918,7 +4028,8 @@ func TestPoolProvisionMaxAttemptsCleanupOnlyClaimsStaleProvisioning(t *testing.T !ok { t.Fatalf("claim exhausted stale provisioning machine deletion ok=%v err=%v", ok, err) } else if deleting.Machine.LifecycleState != "deleting" { - t.Fatalf("exhausted stale provisioning delete claim state = %s, want deleting", deleting.Machine.LifecycleState) + t.Fatalf("exhausted stale provisioning delete claim state = %s, want deleting", + deleting.Machine.LifecycleState) } else if deleting.Machine.LifecycleReasonCode != "provisioning_stale_cleanup" { t.Fatalf( "exhausted stale provisioning delete reason = %s, want provisioning_stale_cleanup", @@ -4103,8 +4214,10 @@ func TestPoolProvisioningAttemptFenceRejectsStaleCompletion(t *testing.T) { if err != nil { t.Fatalf("get current machine: %v", err) } - if current.LifecycleState != "provisioning" || current.ProvisionAttempts != secondClaim.Machine.ProvisionAttempts || - current.ProviderResourceID != "" || current.ProviderProvisionAttemptedAt == nil { + if current.LifecycleState != "provisioning" || + current.ProvisionAttempts != secondClaim.Machine.ProvisionAttempts || + current.ProviderResourceID != "" || + current.ProviderProvisionAttemptedAt == nil { t.Fatalf("stale attempt changed machine: %+v", current) } seedProviderRuntimeMismatchForTest(t, ctx, pool, result.MachineBindings[0].MachineID) @@ -4472,7 +4585,10 @@ func TestOfflinePoolMachineWithoutRuntimeHistoryMovesToCleanupQueueAfterBootstra UPDATE machines SET lifecycle_changed_at = statement_timestamp() - $2::bigint * interval '1 second' WHERE org_id = $1 AND id = $3 - `, testOrgID, int64(executionstore.StaleMachineBootstrapAge/time.Second)+1, result.MachineBindings[0].MachineID); err != nil { + `, + testOrgID, + int64(executionstore.StaleMachineBootstrapAge/time.Second)+1, + result.MachineBindings[0].MachineID); err != nil { t.Fatalf("age offline machine past bootstrap timeout: %v", err) } @@ -4730,7 +4846,8 @@ func TestSystemBootstrapTokenRetryDoesNotRevokePriorToken(t *testing.T) { if err != nil { t.Fatalf("create second system token: %v", err) } - tokens, err := store.Execution().ListAllMachineDaemonTokens(ctx, testOrgID, result.MachineBindings[0].MachineID) + tokens, err := store.Execution(). + ListAllMachineDaemonTokens(ctx, testOrgID, result.MachineBindings[0].MachineID) if err != nil { t.Fatalf("list machine daemon tokens: %v", err) } @@ -4814,7 +4931,8 @@ func TestSystemBootstrapTokenRetryDoesNotRevokePriorToken(t *testing.T) { ); err != nil { t.Fatalf("register daemon runtime with second token: %v", err) } - tokens, err = store.Execution().ListAllMachineDaemonTokens(ctx, testOrgID, result.MachineBindings[0].MachineID) + tokens, err = store.Execution(). + ListAllMachineDaemonTokens(ctx, testOrgID, result.MachineBindings[0].MachineID) if err != nil { t.Fatalf("list machine daemon tokens after registration: %v", err) } @@ -4963,161 +5081,6 @@ tools: } } -func TestClaimNormalModelCallValidatesActiveAgentConfigAtWatermark(t *testing.T) { - t.Parallel() - ctx := context.Background() - pool := openIntegrationDB(t, ctx) - seedMigratedDB(t, ctx, pool) - - store := newIntegrationStore(pool, storage.WithMachinePoolProviders(mergingMachinePoolProviders{})) - user := mustCreateProjectDeveloperUser( - t, - ctx, - store, - "model-context-config@example.com", - "Model Context Config User", - ) - profile := mustCreateConfigAndProfileBookmarkFromYAML( - t, - ctx, - store, - "model-context-config", - "Model Context Config Agent", - ` -instruction: Keep config watermarks typed. -model: - provider_config: openai-prod - name: model-context-config -`, - ) - launch, err := store.Execution().LaunchAgent( - ctx, - executionstore.LaunchAgentInput{ - ProjectID: testProjectID, - ProfileID: profile.ID, - AgentConfigID: profile.CurrentConfigID, - LaunchedBy: userPrincipal(user.ID), - IdempotencyKey: "idem-model-context-config-launch", - }, - ) - if err != nil { - t.Fatalf("launch model-context-config agent: %v", err) - } - input, _, _, err := store.Execution().CreateAgentContentInput( - ctx, - executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - Actor: mustOmnaraActorParams(t, user.ID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"build context"}]`), - IdempotencyKey: "model-context-config-input", - }, - ) - if err != nil { - t.Fatalf("create content input: %v", err) - } - claim, found, err := store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) - if err != nil { - t.Fatalf("claim input work: %v", err) - } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { - t.Fatalf( - "claim input found=%v kind=%v input_ids=%v want %s", - found, - claim.Kind, - claim.Model.InputIDs, - input.ID, - ) - } - lock := claim.RuntimeLock - admitted := claim.Model.AdmittedInputTurn - config := launch.AgentConfig - snapshot, err := store.Execution().CaptureAgentConfigForEventWatermark( - ctx, - testProjectID, - launch.Agent.ID, - admitted.Events[0].Sequence, - ) - if err != nil { - t.Fatalf("capture active config at admitted frontier: %v", err) - } - if snapshot.AgentConfig.ID != config.ID { - t.Fatalf("config at frontier = %s, want %s", snapshot.AgentConfig.ID, config.ID) - } - newConfig := mustCreateAgentConfigFromYAML(t, ctx, store, ` -instruction: Keep changed config typed. -model: - provider_config: openai-prod - name: model-context-config -`) - changed, err := store.Execution().ChangeAgentConfig(ctx, executionstore.ChangeAgentConfigInput{ - CreateAgentConfigInput: changeInputFromRecord(newConfig), - AgentID: launch.Agent.ID, - ActorType: identitystore.PrincipalTypeUser, - ActorID: user.ID, - IdempotencyKey: "idem-model-context-config-change", - }) - if err != nil { - t.Fatalf("change agent config: %v", err) - } - if _, err := store.Execution().ClaimNormalModelCall( - ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: snapshot.InputEventSequence, - }, - ); !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("stale model context frontier error = %v, want ErrAgentNotAdvanceable", err) - } - var staleContexts int - if err := pool.QueryRow(ctx, ` - SELECT count(*) - FROM model_call_contexts - WHERE project_id = $1 AND agent_id = $2 AND input_event_sequence = $3`, - testProjectID, - launch.Agent.ID, - snapshot.InputEventSequence, - ).Scan(&staleContexts); err != nil { - t.Fatalf("count stale model contexts: %v", err) - } - if staleContexts != 0 { - t.Fatalf("stale model contexts = %d, want none", staleContexts) - } - if _, err := store.Execution().ClaimNormalModelCall( - ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: config.ID, - InputEventSequence: changed.ConfigChange.Event.Sequence, - }, - ); !errors.Is( - err, - storeerr.ErrStateTransitionConflict, - ) { - t.Fatalf("stale model context config error = %v, want ErrStateTransitionConflict", err) - } - if _, err := store.Execution().ClaimNormalModelCall( - ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: changed.AgentConfig.ID, - InputEventSequence: changed.ConfigChange.Event.Sequence, - }, - ); err != nil { - t.Fatalf("current model context config error = %v", err) - } -} - func TestModelCallRetryUsesCurrentConfiguredModelRevision(t *testing.T) { t.Parallel() ctx := context.Background() @@ -5181,7 +5144,7 @@ model: if err != nil { t.Fatalf("acquire runtime lock: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, store, @@ -5192,21 +5155,22 @@ model: if !found { t.Fatal("expected input admission") } - configuredModel, err := store.Models().GetConfiguredModel(ctx, testOrgID, launch.AgentConfig.ConfiguredModelID) + configuredModel, err := store.Models(). + GetConfiguredModel(ctx, testOrgID, launch.AgentConfig.ConfiguredModelID) if err != nil { t.Fatalf("load configured model: %v", err) } - initialClaim, err := store.Execution().ClaimNormalModelCall( + prepared1, err := store.Execution().PrepareNormalModelCall( ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: launch.AgentConfig.ID, - InputEventSequence: admitted.Events[0].Sequence, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: launch.Agent.ID, + RuntimeLockID: lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, }, ) + initialClaim := prepared1.Claim + if err != nil { t.Fatalf("claim initial model context: %v", err) } @@ -5245,12 +5209,13 @@ model: if updated.CurrentRevisionID == configuredModel.CurrentRevisionID { t.Fatal("configured model patch did not create a new revision") } - retryClaim, err := store.Execution().ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - PredecessorModelCallContextID: initialClaim.Context.ID, - RuntimeLockID: lock.ID, - }) + retryClaim, err := store.Execution(). + ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ + ProjectID: testProjectID, + AgentID: launch.Agent.ID, + PredecessorModelCallContextID: initialClaim.Context.ID, + RuntimeLockID: lock.ID, + }) if err != nil { t.Fatalf("claim same-frontier retry: %v", err) } @@ -5265,7 +5230,7 @@ model: } if err := executionstore.IntegrationValidateResponseEnvelopeForModelCallContext( ctx, - store.q, + store.pool, modelenvelope.ResponseEnvelope{RequestedProviderModelSlug: updated.ProviderModelSlug}, retryClaim.Context, ); err != nil { @@ -5273,7 +5238,7 @@ model: } if err := executionstore.IntegrationValidateResponseEnvelopeForModelCallContext( ctx, - store.q, + store.pool, modelenvelope.ResponseEnvelope{RequestedProviderModelSlug: configuredModel.ProviderModelSlug}, retryClaim.Context, ); err == nil { @@ -5289,17 +5254,21 @@ model: }); err != nil { t.Fatalf("cancel retrying frontier: %v", err) } - if err := store.Execution().ReleaseAgentRuntimeLock(ctx, testProjectID, launch.Agent.ID, lock.ID); err != nil { + if err := store.Execution().ReleaseAgentRuntimeLock(ctx, + testProjectID, + launch.Agent.ID, + lock.ID); err != nil { t.Fatalf("release canceled runtime: %v", err) } - nextInput, _, _, err := store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - Actor: mustOmnaraActorParams(t, user.ID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"build a new frontier"}]`), - IdempotencyKey: "model-context-revision-new-frontier", - }) + nextInput, _, _, err := store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: launch.Agent.ID, + Actor: mustOmnaraActorParams(t, user.ID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"build a new frontier"}]`), + IdempotencyKey: "model-context-revision-new-frontier", + }) if err != nil { t.Fatalf("create new-frontier input: %v", err) } @@ -5313,7 +5282,7 @@ model: if err != nil { t.Fatalf("acquire new-frontier runtime: %v", err) } - nextAdmitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found = admitNextAgentInputAndOpenTurnForTest( t, ctx, store, @@ -5324,14 +5293,15 @@ model: if !found { t.Fatal("expected new-frontier input admission") } - newFrontierClaim, err := store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: nextLock.ID, - OpeningInputIDs: []uuid.UUID{nextInput.ID}, - AgentConfigID: launch.AgentConfig.ID, - InputEventSequence: nextAdmitted.Events[0].Sequence, - }) + prepared2, err := store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: launch.Agent.ID, + RuntimeLockID: nextLock.ID, + OpeningInputIDs: []uuid.UUID{nextInput.ID}, + }) + newFrontierClaim := prepared2.Claim + if err != nil { t.Fatalf("claim model context at new frontier: %v", err) } @@ -5346,7 +5316,7 @@ model: } if err := executionstore.IntegrationValidateResponseEnvelopeForModelCallContext( ctx, - store.q, + store.pool, modelenvelope.ResponseEnvelope{RequestedProviderModelSlug: updated.ProviderModelSlug}, newFrontierClaim.Context, ); err != nil { @@ -5417,7 +5387,7 @@ model: if err != nil { t.Fatalf("acquire runtime lock: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, store, @@ -5441,17 +5411,17 @@ model: if _, err := store.Models().DeleteProjectModelGrant(ctx, testOrgID, testProjectID, grant.ID); err != nil { t.Fatalf("revoke model grant: %v", err) } - claim, err := store.Execution().ClaimNormalModelCall( + prepared3, err := store.Execution().PrepareNormalModelCall( ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: launch.Agent.ID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: config.ID, - InputEventSequence: admitted.Events[0].Sequence, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: launch.Agent.ID, + RuntimeLockID: lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, }, ) + claim := prepared3.Claim + if err != nil { t.Fatalf("claim context after grant revocation: %v", err) } @@ -5485,7 +5455,8 @@ model: t.Fatalf("expected initial generation 1, got %d", profile.CurrentGeneration) } - retargetInput := func(sourceYAML string, expectedCurrentConfigID uuid.UUID) executionstore.RetargetAgentProfileInput { + retargetInput := func(sourceYAML string, + expectedCurrentConfigID uuid.UUID) executionstore.RetargetAgentProfileInput { config := mustCreateAgentConfigFromYAML( t, ctx, @@ -5519,7 +5490,11 @@ model: retargetedConfigID := updated.CurrentConfigID generation1, err := store.q.GetAgentProfileVersionByGeneration( ctx, - dbsqlc.GetAgentProfileVersionByGenerationParams{ProjectID: testProjectID, ProfileID: profile.ID, Generation: 1}, + dbsqlc.GetAgentProfileVersionByGenerationParams{ + ProjectID: testProjectID, + ProfileID: profile.ID, + Generation: 1, + }, ) if err != nil { t.Fatalf("load profile generation 1: %v", err) @@ -5529,7 +5504,11 @@ model: } generation2, err := store.q.GetAgentProfileVersionByGeneration( ctx, - dbsqlc.GetAgentProfileVersionByGenerationParams{ProjectID: testProjectID, ProfileID: profile.ID, Generation: 2}, + dbsqlc.GetAgentProfileVersionByGenerationParams{ + ProjectID: testProjectID, + ProfileID: profile.ID, + Generation: 2, + }, ) if err != nil { t.Fatalf("load profile generation 2: %v", err) @@ -5658,7 +5637,10 @@ model: t.Fatalf("launch config-only agent: %v", err) } if configOnly.Agent.CurrentConfigID != retargetedConfigID || configOnly.Agent.Name != "" { - t.Fatalf("config-only launch should use requested config and have no name, got agent=%+v", configOnly.Agent) + t.Fatalf( + "config-only launch should use requested config and have no name, got agent=%+v", + configOnly.Agent, + ) } configOnlyName := "Config Only Named" @@ -5677,7 +5659,10 @@ model: t.Fatalf("launch named config-only agent: %v", err) } if configOnlyNamed.Agent.Name != "Config Only Named" { - t.Fatalf("config-only launch with explicit name should set agent name, got agent=%+v", configOnlyNamed.Agent) + t.Fatalf( + "config-only launch with explicit name should set agent name, got agent=%+v", + configOnlyNamed.Agent, + ) } other := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "retarget-other", "Other Agent", ` diff --git a/internal/storage/executionstore/agent_launch_machines.go b/internal/storage/executionstore/agent_launch_machines.go index 03a191fcf..5d54e50d4 100644 --- a/internal/storage/executionstore/agent_launch_machines.go +++ b/internal/storage/executionstore/agent_launch_machines.go @@ -74,7 +74,7 @@ func expandLaunchMachineBindingRequests( func (s *Store) resolveLaunchMachineSourcesTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, qtx *dbsqlc.Queries, orgID, projectID uuid.UUID, sources []launchMachineSource, @@ -90,7 +90,7 @@ func (s *Store) resolveLaunchMachineSourcesTx( func (s *Store) resolveLaunchPoolMachineSourcesTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, qtx *dbsqlc.Queries, orgID, projectID uuid.UUID, sources []launchMachineSource, @@ -151,7 +151,7 @@ func (s *Store) resolveLaunchPoolMachineSourcesTx( func lockLaunchMachineSourcesTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, orgID uuid.UUID, sources []launchMachineSource, additionalMachineIDs []uuid.UUID, diff --git a/internal/storage/executionstore/agent_launch_store.go b/internal/storage/executionstore/agent_launch_store.go index 20bcf713e..47356ba07 100644 --- a/internal/storage/executionstore/agent_launch_store.go +++ b/internal/storage/executionstore/agent_launch_store.go @@ -10,11 +10,11 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/dbsafe" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/resourcename" "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -105,13 +105,14 @@ func (s *Store) launchAgentOnce( ctx context.Context, input LaunchAgentInput, ) (LaunchAgentResult, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return LaunchAgentResult{}, fmt.Errorf("begin launch agent: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - result, err := s.launchAgentTx(ctx, tx, s.q.WithTx(tx), txNotifications, input) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + result, err := s.launchAgentTx(ctx, unit, dbsqlc.New(tx), input) if err != nil { return LaunchAgentResult{}, err } @@ -119,7 +120,7 @@ func (s *Store) launchAgentOnce( if !result.Created { scope = "idempotent launch agent" } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, scope); err != nil { + if err := unit.Commit(ctx, scope); err != nil { return LaunchAgentResult{}, err } return result, nil @@ -127,11 +128,11 @@ func (s *Store) launchAgentOnce( func (s *Store) launchAgentTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, - txNotifications *notifications.TxNotifications, input LaunchAgentInput, ) (LaunchAgentResult, error) { + tx := unit.DB() project, err := loadProjectTx(ctx, qtx, input.ProjectID) if err != nil { return LaunchAgentResult{}, err @@ -265,40 +266,34 @@ func (s *Store) launchAgentTx( if input.admission != nil { insertInput.ID = input.admission.AgentID } + if err := lockResourceCreation(ctx, qtx, resourceAgents, input.ProjectID.String()); err != nil { + return LaunchAgentResult{}, err + } + var sharedBindings []dbsqlc.ListParentMachineBindingsForSharingRow if input.Subagent != nil { - if err := admitSubagentLaunchTx(ctx, tx, qtx, input.ProjectID, *input.Subagent); err != nil { + machineSources = nil + sharedBindings, err = lockParentMachineBindingsForSharingTx( + ctx, tx, qtx, project.OrgID, input.ProjectID, input.Subagent.ParentAgentID, + ) + if err != nil { + return LaunchAgentResult{}, err + } + if err := admitSubagentLaunchTx(ctx, unit, qtx, input.ProjectID, *input.Subagent); err != nil { return LaunchAgentResult{}, err } insertInput.ParentAgentID = input.Subagent.ParentAgentID insertInput.SubagentKey = input.Subagent.Key - machineSources = nil } - agent, inserted, err := insertAdmittedAgentTx(ctx, tx, qtx, insertInput) + if err := s.resolveLaunchMachineSourcesTx(ctx, tx, qtx, project.OrgID, input.ProjectID, machineSources); err != nil { + return LaunchAgentResult{}, err + } + agent, inserted, err := insertAdmittedAgentTx(ctx, unit, qtx, insertInput) if err != nil { return LaunchAgentResult{}, err } if !inserted { return LaunchAgentResult{Agent: agent}, nil } - if err := s.resolveLaunchMachineSourcesTx( - ctx, - tx, - qtx, - project.OrgID, - input.ProjectID, - machineSources, - ); err != nil { - return LaunchAgentResult{}, err - } - var sharedBindings []dbsqlc.ListParentMachineBindingsForSharingRow - if input.Subagent != nil { - sharedBindings, err = lockParentMachineBindingsForSharingTx( - ctx, tx, qtx, project.OrgID, input.ProjectID, input.Subagent.ParentAgentID, - ) - if err != nil { - return LaunchAgentResult{}, err - } - } result := LaunchAgentResult{ Agent: agent, AgentConfig: config, @@ -321,7 +316,7 @@ func (s *Store) launchAgentTx( return LaunchAgentResult{}, err } } - configChange, err := activateNewAgentConfigTx(ctx, txNotifications, tx, qtx, ActivateAgentConfigInput{ + configChange, err := activateNewAgentConfigTx(ctx, unit, qtx, ActivateAgentConfigInput{ ProjectID: input.ProjectID, AgentID: agent.ID, AgentConfigID: config.ID, @@ -409,8 +404,7 @@ func (s *Store) launchAgentTx( if initial != nil { if err := s.insertLaunchInitialContentInputTx( ctx, - tx, - txNotifications, + unit, agent, input, *initial, diff --git a/internal/storage/executionstore/agent_machine_source_reconciliation.go b/internal/storage/executionstore/agent_machine_source_reconciliation.go index 1a25881f8..59043cb67 100644 --- a/internal/storage/executionstore/agent_machine_source_reconciliation.go +++ b/internal/storage/executionstore/agent_machine_source_reconciliation.go @@ -11,6 +11,7 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -19,7 +20,7 @@ import ( func (s *Store) reconcileAgentMachineSourcesTx( ctx context.Context, txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, projectID, agentID uuid.UUID, currentContract, nextContract agentconfig.RuntimeContract, @@ -68,7 +69,7 @@ func (s *Store) reconcileAgentMachineSourcesTx( if err := completeExecutionRevokedProcessesTx( ctx, txNotifications, - tx, + unit, qtx, executionRevokedProcessScope{ projectID: projectID, diff --git a/internal/storage/executionstore/agent_rows.go b/internal/storage/executionstore/agent_rows.go index 3c34098dd..51ad21264 100644 --- a/internal/storage/executionstore/agent_rows.go +++ b/internal/storage/executionstore/agent_rows.go @@ -9,31 +9,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" ) -func agentRecordFromInsertSQLC(row dbsqlc.InsertAgentRow) AgentRecord { - record := agentRecordFromSQLC( - row.ID, - row.OrgID, - row.ProjectID, - row.State, - row.Name, - row.AgentProfileID, - row.CurrentConfigID, - row.InteractionTargetID, - row.IdempotencyKey, - row.NextEventSequence, - row.CreatedAt, - row.UpdatedAt, - row.ArchivedAt, - row.ParentAgentID, - row.SubagentKey, - ) - record.Model = AgentModelDisplay{ - ProviderConfig: row.ModelProviderConfigName, - Name: row.ModelName, - } - return record -} - func agentRecordFromIdempotencySQLC(row dbsqlc.GetAgentByIdempotencyKeyRow) AgentRecord { record := agentRecordFromSQLC( row.ID, diff --git a/internal/storage/executionstore/agent_runtime_cancellation_store.go b/internal/storage/executionstore/agent_runtime_cancellation_store.go index 7951b803f..e5c5bc059 100644 --- a/internal/storage/executionstore/agent_runtime_cancellation_store.go +++ b/internal/storage/executionstore/agent_runtime_cancellation_store.go @@ -2,18 +2,11 @@ package executionstore import ( "context" - "encoding/json" - "errors" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/events" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" ) type CancelAgentResult struct { @@ -31,419 +24,24 @@ type CancelAgentInput struct { const cancelReasonAgentCanceled = "agent_canceled" -func (s *Store) CancelAgent( - ctx context.Context, - input CancelAgentInput, -) (CancelAgentResult, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil { - return CancelAgentResult{}, errors.New("project and agent are required") - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) +func (s *Store) CancelAgent(ctx context.Context, input CancelAgentInput) (CancelAgentResult, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return CancelAgentResult{}, fmt.Errorf("begin cancel agent: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if err := lifecyclelock.Agents(ctx, tx, []lifecyclelock.AgentRef{{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }}); err != nil { return CancelAgentResult{}, err } - actorID, err := resolveActorTx(ctx, qtx, input.ProjectID, input.Actor) + defer func() { _ = unit.Rollback(ctx) }() + actor, err := resolveActorTx(ctx, dbsqlc.New(unit.DB()), input.ProjectID, input.Actor) if err != nil { return CancelAgentResult{}, err } - result, err := cancelAgentTx( - ctx, - txNotifications, - tx, - qtx, - cancelAgentTxInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ActorID: actorID, - ReasonCode: cancelReasonAgentCanceled, - ModelCallMessage: "The model call was canceled by an explicit agent cancellation.", - }, - ) + result, err := h.Cancel(ctx, agentexecution.CancelInput{ActorID: actor, Reason: cancelReasonAgentCanceled, + Message: "The model call was canceled by an explicit agent cancellation."}) if err != nil { return CancelAgentResult{}, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - fmt.Sprintf("cancel agent %s", input.AgentID.String()), - ); err != nil { - return CancelAgentResult{}, err - } - return result, nil -} - -type cancelAgentTxInput struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ActorID uuid.UUID - ReasonCode string - ModelCallMessage string - CancelRuntimeWithoutContinuableTurn bool -} - -func cancelAgentTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - input cancelAgentTxInput, -) (CancelAgentResult, error) { - projectID, agentID, actorID := input.ProjectID, input.AgentID, input.ActorID - latest, err := qtx.LatestAgentEvent(ctx, dbsqlc.LatestAgentEventParams{ProjectID: projectID, AgentID: agentID}) - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return CancelAgentResult{}, fmt.Errorf("load latest agent event for cancel: %w", err) - } - - afterSequence := int64(0) - if latest.ID != uuid.Nil { - afterSequence = latest.Sequence - } - var event events.Event - currentTurn, currentTurnErr := qtx.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: projectID, AgentID: agentID}, - ) - if currentTurnErr != nil && !errors.Is(currentTurnErr, pgx.ErrNoRows) { - return CancelAgentResult{}, fmt.Errorf("load current turn for cancel: %w", currentTurnErr) - } - hasActiveContexts, err := qtx.AgentHasLiveModelCallContexts( - ctx, - dbsqlc.AgentHasLiveModelCallContextsParams{ProjectID: projectID, AgentID: agentID}, - ) - if err != nil { - return CancelAgentResult{}, fmt.Errorf("check active model call contexts for cancel: %w", err) - } - runtimeCancelRequested := false - var runtimeToCancel AgentRuntimeLockRecord - recordRuntimeCancel := func(record AgentRuntimeLockRecord) { - if record.ID == uuid.Nil { - return - } - runtimeCancelRequested = true - runtimeToCancel = record - txNotifications.AddWorkerControlCancel(record.WorkerProcessID, agentID, record.ID) - } - recordPendingRuntimeCancel := func() error { - row, err := qtx.GetPendingAgentRuntimeCancel(ctx, dbsqlc.GetPendingAgentRuntimeCancelParams{ - ProjectID: projectID, - AgentID: agentID, - }) - if errors.Is(err, pgx.ErrNoRows) { - return nil - } - if err != nil { - return fmt.Errorf("load pending runtime cancel: %w", err) - } - recordRuntimeCancel(agentRuntimeLockRecordFromSQLC(row)) - return nil - } - affectedTurnID := uuid.Nil - if currentTurnErr == nil { - affectedTurnID = currentTurn.ID - } - if currentTurnErr == nil || hasActiveContexts || input.CancelRuntimeWithoutContinuableTurn { - params := dbsqlc.RequestAgentRuntimeCancelParams{ - ProjectID: projectID, - AgentID: agentID, - } - if row, err := qtx.RequestAgentRuntimeCancel(ctx, params); err == nil { - recordRuntimeCancel(agentRuntimeLockRecordFromSQLC(row)) - } else if errors.Is(err, pgx.ErrNoRows) { - if err := recordPendingRuntimeCancel(); err != nil { - return CancelAgentResult{}, err - } - } else { - return CancelAgentResult{}, fmt.Errorf("request runtime cancel: %w", err) - } - } else if err := recordPendingRuntimeCancel(); err != nil { + if err = unit.Commit(ctx, "cancel agent"); err != nil { return CancelAgentResult{}, err } - if err := terminalizeAgentModelCallsForLifecycleUnderAgentLockTx( - ctx, - qtx, - projectID, - agentID, - runtimeToCancel.ID, - input.ReasonCode, - input.ModelCallMessage, - ); err != nil { - return CancelAgentResult{}, err - } - affected := affectedTurnID != uuid.Nil || hasActiveContexts - cancelInputID := uuid.Nil - if affected { - cancelSteeringParams := dbsqlc.CancelSteeringAgentInputsForAgentParams{ - ProjectID: projectID, - AgentID: agentID, - } - if _, err := qtx.CancelSteeringAgentInputsForAgent(ctx, cancelSteeringParams); err != nil { - return CancelAgentResult{}, fmt.Errorf("cancel steering agent inputs: %w", err) - } - controlType := "cancel_current" - idempotencyKey := fmt.Sprintf("agent-cancel:%s:after:%d", agentID.String(), afterSequence) - controlInput, insertErr := qtx.InsertControlAgentInput(ctx, dbsqlc.InsertControlAgentInputParams{ - ProjectID: projectID, - AgentID: agentID, - ActorID: storeutil.IDFromNil(actorID), - ControlType: &controlType, - IdempotencyScope: storeutil.TextFromEmpty("agent_control"), - InputIdempotencyKey: storeutil.TextFromEmpty(idempotencyKey), - Metadata: json.RawMessage(`{}`), - }) - if insertErr != nil { - return CancelAgentResult{}, fmt.Errorf("insert cancel control input: %w", insertErr) - } - inputRecord := agentInputRecordFromControlSQLC(controlInput) - cancelInputID = inputRecord.ID - eventRecord, eventErr := appendTypedAgentEventTx(ctx, txNotifications, tx, AppendTypedAgentEventInput{ - ProjectID: projectID, - AgentID: agentID, - TurnID: affectedTurnID, - Kind: events.KindAgentInput, - IdempotencyKey: "agent_input:" + inputRecord.ID.String(), - AgentInputID: inputRecord.ID, - }) - if eventErr != nil { - return CancelAgentResult{}, eventErr - } - if err := updateAgentTurnLatestEventQuery( - ctx, - qtx, - projectID, - agentID, - affectedTurnID, - eventRecord.Event.ID, - uuid.Nil, - ); err != nil { - return CancelAgentResult{}, err - } - resolveParams := dbsqlc.ResolveControlAgentInputParams{ - ProjectID: projectID, - AgentID: agentID, - ID: inputRecord.ID, - ControlType: &controlType, - EventID: &eventRecord.Event.ID, - } - if err := qtx.ResolveControlAgentInput(ctx, resolveParams); err != nil { - return CancelAgentResult{}, fmt.Errorf("resolve cancel control input: %w", err) - } - event = eventRecord.Event - } - var interactionRows []dbsqlc.AgentInteractionReadProjection - if currentTurnErr == nil { - interactionIDs, cancelErr := qtx.CancelOpenAgentInteractionsForAgent( - ctx, - dbsqlc.CancelOpenAgentInteractionsForAgentParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: currentTurn.ID, - Reason: input.ReasonCode, - ResolvedByInputID: storeutil.IDFromNil(cancelInputID), - }, - ) - if cancelErr != nil { - return CancelAgentResult{}, fmt.Errorf("cancel open agent interactions: %w", cancelErr) - } - interactionRows, err = qtx.ListAgentInteractionsByIDs( - ctx, - dbsqlc.ListAgentInteractionsByIDsParams{ - ProjectID: projectID, - AgentID: agentID, - Ids: interactionIDs, - }, - ) - if err != nil { - return CancelAgentResult{}, fmt.Errorf("load canceled agent interactions: %w", err) - } - } - interactionUpdates := make(map[uuid.UUID]*notifications.AgentInteractionUpdate, len(interactionRows)) - for _, row := range interactionRows { - interaction := agentInteractionRecordFromSQLC(row) - interactionUpdates[interaction.ToolCallID] = interactionUpdateFromRecord(interaction) - if interaction.TurnID != uuid.Nil { - params := dbsqlc.MarkAgentWakeupParams{ - ProjectID: interaction.ProjectID, - AgentID: interaction.AgentID, - Metadata: []byte(`{"reason":"interaction_canceled"}`), - } - if err := qtx.MarkAgentWakeup(ctx, params); err != nil { - return CancelAgentResult{}, err - } - } - } - var toolCallRows []dbsqlc.CancelNonTerminalToolCallsForAgentRow - if currentTurnErr == nil { - processRows, err := qtx.CancelUnresolvedProcessesForAgentTurn( - ctx, - dbsqlc.CancelUnresolvedProcessesForAgentTurnParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: currentTurn.ID, - }, - ) - if err != nil { - return CancelAgentResult{}, fmt.Errorf( - "cancel unresolved processes with their agent turn: %w", - err, - ) - } - if _, err := qtx.CancelQueuedProcessActionsForAgentTurn( - ctx, - dbsqlc.CancelQueuedProcessActionsForAgentTurnParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: currentTurn.ID, - }, - ); err != nil { - return CancelAgentResult{}, fmt.Errorf( - "cancel queued process actions before grant: %w", - err, - ) - } - if _, err := qtx.CancelAcceptedProcessActionsForAgentTurn( - ctx, - dbsqlc.CancelAcceptedProcessActionsForAgentTurnParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: currentTurn.ID, - }, - ); err != nil { - return CancelAgentResult{}, fmt.Errorf( - "mark accepted process actions unknown after cancel: %w", - err, - ) - } - for _, process := range processRows { - if err := completeQueuedProcessActionsFailedTx( - ctx, - txNotifications, - tx, - qtx, - process.OrgID, - process.ID, - "agent_canceled_before_grant", - ); err != nil { - return CancelAgentResult{}, err - } - if err := completeAcceptedProcessActionsWithoutEvidenceTx( - ctx, - txNotifications, - tx, - qtx, - process.OrgID, - process.ID, - "agent_canceled_after_grant", - ); err != nil { - return CancelAgentResult{}, err - } - if ProcessState(process.State) == ProcessStateUnknown { - txNotifications.AddDaemonProcessTermination( - process.MachineID, - process.ID, - ) - } - } - toolCallRows, err = qtx.CancelNonTerminalToolCallsForAgent( - ctx, - dbsqlc.CancelNonTerminalToolCallsForAgentParams{ - ProjectID: projectID, - AgentID: agentID, - TurnID: currentTurn.ID, - }, - ) - if err != nil { - return CancelAgentResult{}, fmt.Errorf("cancel non-terminal tool calls: %w", err) - } - } - for _, row := range toolCallRows { - resultRecord := toolCallRecordFromCancelSQLC(row) - if affectedTurnID != resultRecord.TurnID { - return CancelAgentResult{}, storeerr.ErrStateTransitionConflict - } - } - for _, row := range toolCallRows { - resultRecord := toolCallRecordFromCancelSQLC(row) - resultRecord.ResultContentParts, err = canceledToolResultContentParts() - if err != nil { - return CancelAgentResult{}, fmt.Errorf("marshal canceled tool result content parts: %w", err) - } - if _, err := appendToolResultEventTx( - ctx, txNotifications, tx, resultRecord, interactionUpdates[resultRecord.ID], - ); err != nil { - return CancelAgentResult{}, fmt.Errorf("append canceled tool result event: %w", err) - } - } - wakeupMetadata, err := marshalJSON(map[string]string{"reason": input.ReasonCode}) - if err != nil { - return CancelAgentResult{}, fmt.Errorf("marshal canceled agent wakeup metadata: %w", err) - } - params := dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - Metadata: wakeupMetadata, - } - if err := qtx.ReconcileAgentWakeup(ctx, params); err != nil { - return CancelAgentResult{}, fmt.Errorf("reconcile canceled agent wakeup: %w", err) - } - return CancelAgentResult{ - Event: event, - RuntimeCancelRequested: runtimeCancelRequested, - Affected: affected, - ActorID: actorID, - }, nil -} - -func terminalizeAgentModelCallsForLifecycleUnderAgentLockTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeLockID uuid.UUID, - reasonCode, contextMessage string, -) error { - if runtimeLockID != uuid.Nil { - errorDetails, err := marshalJSON(map[string]any{ - "code": reasonCode, - "message": contextMessage, - "runtime_lock_id": runtimeLockID, - }) - if err != nil { - return fmt.Errorf("marshal lifecycle model call details: %w", err) - } - if _, err := qtx.CancelRuntimeModelCallContextsForLifecycle( - ctx, - dbsqlc.CancelRuntimeModelCallContextsForLifecycleParams{ - ErrorCode: reasonCode, - ErrorMessage: contextMessage, - ErrorDetails: errorDetails, - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtimeLockID, - }, - ); err != nil { - return fmt.Errorf("cancel lifecycle model call contexts: %w", err) - } - } - live, err := qtx.AgentHasLiveModelCallContexts( - ctx, - dbsqlc.AgentHasLiveModelCallContextsParams{ - ProjectID: projectID, - AgentID: agentID, - }, - ) - if err != nil { - return fmt.Errorf("verify lifecycle model call contexts: %w", err) - } - if live { - return storeerr.ErrStateTransitionConflict - } - return nil + return CancelAgentResult{Event: executionEvent(input.AgentID, result.Event, events.KindAgentInput), + Affected: result.Changed, RuntimeCancelRequested: result.RuntimeID != uuid.Nil, ActorID: actor}, nil } diff --git a/internal/storage/executionstore/agent_runtime_lock_reaping_store.go b/internal/storage/executionstore/agent_runtime_lock_reaping_store.go index 0d1fc65a2..cf12a279d 100644 --- a/internal/storage/executionstore/agent_runtime_lock_reaping_store.go +++ b/internal/storage/executionstore/agent_runtime_lock_reaping_store.go @@ -4,12 +4,10 @@ import ( "context" "errors" "fmt" - "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/errutil" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -78,25 +76,17 @@ func (s *Store) reapExpiredAgentRuntimeLock( ctx context.Context, projectID, agentID, runtimeLockID uuid.UUID, ) (bool, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return false, fmt.Errorf("begin reap expired agent runtime lock: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - reaped, err := reapExpiredAgentRuntimeLockTx( - ctx, - txNotifications, - tx, - projectID, - agentID, - runtimeLockID, - s.modelCallRetryDelay, - ) + defer func() { _ = unit.Rollback(ctx) }() + + reaped, err := reapExpiredAgentRuntimeLockTx(ctx, unit, projectID, agentID, runtimeLockID) if err != nil || !reaped { return false, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "reap expired agent runtime lock"); err != nil { + if err := unit.Commit(ctx, "reap expired agent runtime lock"); err != nil { return false, err } return true, nil @@ -104,98 +94,19 @@ func (s *Store) reapExpiredAgentRuntimeLock( func reapExpiredAgentRuntimeLockTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, projectID, agentID, runtimeLockID uuid.UUID, - retryBackoff func(int, string) time.Duration, ) (bool, error) { - qtx := dbsqlc.New(tx) - lockedAgent, err := tryLockAgentWithParentTx(ctx, tx, qtx, projectID, agentID) + locked, err := tryLockAgentWithParentTx(ctx, unit, dbsqlc.New(unit.DB()), projectID, agentID) if errors.Is(err, storeerr.ErrNotFound) { return false, nil } - if err != nil { - return false, fmt.Errorf("lock agent for expired runtime lock reap: %w", err) - } - if !lockedAgent { - return false, nil - } - locked, err := qtx.LockExpiredAgentRuntimeLockForReap( - ctx, - dbsqlc.LockExpiredAgentRuntimeLockForReapParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeLockID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return false, nil - } - if err != nil { - return false, fmt.Errorf("lock expired agent runtime lock for reap: %w", err) - } - if err := recoverRuntimeModelCallContextTx( - ctx, - txNotifications, - tx, - qtx, - locked.ProjectID, - locked.AgentID, - locked.ID, - runtimeModelCallRecoveryEvidence{ - Code: "runtime_lease_expired_before_model_result_acceptance", - Message: "runtime lease expired before the model result was durably accepted", - }, - retryBackoff, - ); err != nil { + if err != nil || !locked { return false, err } - if err := failRuntimeToolCallsTx( - ctx, - txNotifications, - tx, - locked.ProjectID, - locked.AgentID, - locked.ID, - "runtime_lock_stale", - ); err != nil { - return false, err - } - if err := failQueuedRuntimeWorkTx( - ctx, - txNotifications, - tx, - qtx, - locked.ProjectID, - locked.AgentID, - locked.ID, - "runtime_lock_stale", - ); err != nil { - return false, err - } - changed, err := qtx.DeleteAgentRuntimeLockForReap( - ctx, - dbsqlc.DeleteAgentRuntimeLockForReapParams{ - ProjectID: locked.ProjectID, - AgentID: locked.AgentID, - ID: locked.ID, - }, - ) + h, err := unit.Handle(projectID, agentID) if err != nil { - return false, fmt.Errorf("delete expired agent runtime lock: %w", err) - } - if changed != 1 { - return false, storeerr.ErrRuntimeLockInactive - } - if err := qtx.ReconcileAgentWakeup( - ctx, - dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: locked.ProjectID, - AgentID: locked.AgentID, - Metadata: []byte(`{"reason":"runtime_lock_reap"}`), - }, - ); err != nil { - return false, fmt.Errorf("reconcile wakeup after expired runtime lock reap: %w", err) + return false, err } - return true, nil + return h.Reap(ctx, runtimeLockID) } diff --git a/internal/storage/executionstore/agent_runtime_locks_integration_test.go b/internal/storage/executionstore/agent_runtime_locks_integration_test.go index 11a325cb2..19deb2710 100644 --- a/internal/storage/executionstore/agent_runtime_locks_integration_test.go +++ b/internal/storage/executionstore/agent_runtime_locks_integration_test.go @@ -10,13 +10,11 @@ import ( "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/testutil/integrationdb" ) @@ -135,14 +133,11 @@ func TestAgentRuntimeLockMutationsRejectWrongProject(t *testing.T) { ); !errors.Is(err, storeerr.ErrRuntimeLockInactive) { t.Fatalf("wrong-project renewal error = %v, want runtime lock inactive", err) } - if _, err := dbsqlc.New(fixture.Pool).RequestAgentRuntimeCancel( - ctx, - dbsqlc.RequestAgentRuntimeCancelParams{ - ProjectID: wrongProjectID, - AgentID: fixture.AgentID, - }, - ); !errors.Is(err, pgx.ErrNoRows) { - t.Fatalf("wrong-project cancellation error = %v, want no rows", err) + if _, + err := fixture.Store.Execution().CancelAgent(ctx, + executionstore.CancelAgentInput{ProjectID: wrongProjectID, + AgentID: fixture.AgentID}); err == nil { + t.Fatal("wrong-project cancellation succeeded") } if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( ctx, @@ -310,55 +305,11 @@ FOR EACH ROW EXECUTE FUNCTION fail_runtime_lock_reap_commit(); t.Fatalf("load runtime locks after isolated failure: %v", err) } if !failedExists || laterExists { - t.Fatalf("runtime locks after isolated failure: failed=%v later=%v, want true/false", failedExists, laterExists) - } -} - -func TestAgentRuntimeLockRenewalAdvancesLeaseTogether(t *testing.T) { - t.Parallel() - ctx := context.Background() - fixture := newRuntimeLockLeaseFixture(t, ctx) - lock := fixture.acquire(t, ctx, testWorkerProcessID, time.Minute) - if _, err := fixture.Pool.Exec( - ctx, - `UPDATE agent_runtime_locks -SET started_at = statement_timestamp() - interval '5 minutes', - renewed_at = statement_timestamp() - interval '4 minutes', - lease_expires_at = statement_timestamp() - interval '3 minutes' -WHERE id = $1`, - lock.ID, - ); err != nil { - t.Fatalf("age runtime-lock lease: %v", err) - } - var agedRenewedAt, agedLeaseExpiresAt time.Time - if err := fixture.Pool.QueryRow( - ctx, - `SELECT renewed_at, lease_expires_at FROM agent_runtime_locks WHERE id = $1`, - lock.ID, - ).Scan(&agedRenewedAt, &agedLeaseExpiresAt); err != nil { - t.Fatalf("load aged runtime lock: %v", err) - } - - leaseDuration := 41 * time.Second - renewal, err := fixture.Store.Execution().RenewAgentRuntimeLock( - ctx, - testProjectID, - fixture.AgentID, - lock.ID, - leaseDuration, - ) - if err != nil { - t.Fatalf("renew agent runtime lock: %v", err) - } - renewedLock := renewal.RuntimeLock - if !renewedLock.RenewedAt.After(agedRenewedAt) { - t.Fatalf("renewal did not advance: before=%s after=%s", agedRenewedAt, renewedLock.RenewedAt) - } - if !renewedLock.LeaseExpiresAt.After(agedLeaseExpiresAt) { - t.Fatalf("lease deadline did not advance: before=%s after=%s", agedLeaseExpiresAt, renewedLock.LeaseExpiresAt) - } - if got := renewedLock.LeaseExpiresAt.Sub(renewedLock.RenewedAt); got != leaseDuration { - t.Fatalf("renewed lease duration = %s, want %s", got, leaseDuration) + t.Fatalf( + "runtime locks after isolated failure: failed=%v later=%v, want true/false", + failedExists, + laterExists, + ) } } @@ -450,7 +401,11 @@ func TestAgentRuntimeLockAcquisitionWaitGrantsFullLeaseAfterAgentUnlock(t *testi t.Fatalf("acquire after agent unlock: %v", result.err) } if result.lock.RenewedAt.Before(releasedAfter) { - t.Fatalf("lease started before row-lock wait ended: renewed=%s marker=%s", result.lock.RenewedAt, releasedAfter) + t.Fatalf( + "lease started before row-lock wait ended: renewed=%s marker=%s", + result.lock.RenewedAt, + releasedAfter, + ) } if got := result.lock.LeaseExpiresAt.Sub(result.lock.RenewedAt); got != leaseDuration { t.Fatalf("post-wait lease duration = %s, want full %s", got, leaseDuration) @@ -491,7 +446,7 @@ func TestAgentRuntimeLockRenewalWaitGrantsFullLeaseAfterRuntimeUnlock(t *testing ) renewed <- renewalResult{renewal: renewedLock, err: renewErr} }() - integrationdb.WaitForNamedLockWaiters(t, ctx, fixture.Pool, "LockAgentRuntimeLockForRenewal", 1) + integrationdb.WaitForNamedLockWaiters(t, ctx, fixture.Pool, "LockExecutionRuntime", 1) var releasedAfter time.Time if err := fixture.Pool.QueryRow(ctx, `SELECT statement_timestamp()`).Scan(&releasedAfter); err != nil { t.Fatalf("read database time before runtime unlock: %v", err) @@ -514,7 +469,11 @@ func TestAgentRuntimeLockRenewalWaitGrantsFullLeaseAfterRuntimeUnlock(t *testing } renewedLock := result.renewal.RuntimeLock if renewedLock.RenewedAt.Before(releasedAfter) { - t.Fatalf("lease renewed before row-lock wait ended: renewed=%s marker=%s", renewedLock.RenewedAt, releasedAfter) + t.Fatalf( + "lease renewed before row-lock wait ended: renewed=%s marker=%s", + renewedLock.RenewedAt, + releasedAfter, + ) } if got := renewedLock.LeaseExpiresAt.Sub(renewedLock.RenewedAt); got != leaseDuration { t.Fatalf("post-wait renewed lease duration = %s, want full %s", got, leaseDuration) @@ -637,7 +596,11 @@ WHERE id IN ($1, $2)`, t.Fatalf("load runtime locks after contended reap: %v", err) } if !firstExists || secondExists { - t.Fatalf("runtime locks after contended reap: first=%v second=%v, want true/false", firstExists, secondExists) + t.Fatalf( + "runtime locks after contended reap: first=%v second=%v, want true/false", + firstExists, + secondExists, + ) } if err := runtimeTx.Commit(ctx); err != nil { t.Fatalf("release first runtime row: %v", err) @@ -658,14 +621,14 @@ func TestAgentRuntimeLockRenewalWinningReaperRacePreservesLease(t *testing.T) { lock := fixture.acquire(t, ctx, testWorkerProcessID, time.Minute) expireAgentRuntimeLockForTest(t, ctx, fixture.Store, lock.ID) - renewalTx, err := fixture.Pool.Begin(ctx) + renewalTx, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) if err != nil { t.Fatalf("begin renewal transaction: %v", err) } defer func() { _ = renewalTx.Rollback(ctx) }() renewal, err := executionstore.IntegrationRenewAgentRuntimeLockTx( ctx, - dbsqlc.New(renewalTx), + renewalTx, testProjectID, fixture.AgentID, lock.ID, @@ -687,7 +650,7 @@ func TestAgentRuntimeLockRenewalWinningReaperRacePreservesLease(t *testing.T) { if reaped != 0 { t.Fatalf("reaper deleted %d locks while renewal owns runtime row, want 0", reaped) } - if err := renewalTx.Commit(ctx); err != nil { + if err := renewalTx.Commit(ctx, "renewal race"); err != nil { t.Fatalf("commit winning renewal: %v", err) } if err := fixture.Store.Execution().EnsureRuntimeLockActive( @@ -731,20 +694,17 @@ func TestAgentRuntimeLockReaperWinningRaceFencesOldWorker(t *testing.T) { t.Fatalf("expired runtime-lock write before reap = %v, want inactive runtime lock", err) } - reapTx, err := fixture.Pool.BeginTx(ctx, pgx.TxOptions{}) + reapTx, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) if err != nil { t.Fatalf("begin reaper transaction: %v", err) } defer func() { _ = reapTx.Rollback(ctx) }() - txNotifications := notifications.NewTxNotifications() reaped, err := executionstore.IntegrationReapExpiredAgentRuntimeLockTx( ctx, - txNotifications, reapTx, testProjectID, fixture.AgentID, lock.ID, - executionstore.ModelCallRetryBackoff, ) if err != nil { t.Fatalf("reap expired runtime lock in transaction: %v", err) @@ -764,13 +724,8 @@ func TestAgentRuntimeLockReaperWinningRaceFencesOldWorker(t *testing.T) { ) renewalDone <- renewalErr }() - integrationdb.WaitForNamedLockWaiters(t, ctx, fixture.Pool, "LockAgentRuntimeLockForRenewal", 1) - if err := fixture.Store.Execution().IntegrationCommitTxWithNotifications( - ctx, - reapTx, - txNotifications, - "commit winning reaper", - ); err != nil { + integrationdb.WaitForNamedLockWaiters(t, ctx, fixture.Pool, "LockExecutionRuntime", 1) + if err := reapTx.Commit(ctx, "commit winning reaper"); err != nil { t.Fatalf("commit winning reaper: %v", err) } if err := <-renewalDone; !errors.Is(err, storeerr.ErrRuntimeLockInactive) { @@ -802,20 +757,17 @@ func TestConcurrentRuntimeLockReapersApplyRecoveryOnce(t *testing.T) { t.Fatalf("clear agent wakeup before reap: %v", err) } - reapTx, err := fixture.Pool.Begin(ctx) + reapTx, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) if err != nil { t.Fatalf("begin winning reaper: %v", err) } defer func() { _ = reapTx.Rollback(ctx) }() - txNotifications := notifications.NewTxNotifications() won, err := executionstore.IntegrationReapExpiredAgentRuntimeLockTx( ctx, - txNotifications, reapTx, testProjectID, fixture.AgentID, lock.ID, - executionstore.ModelCallRetryBackoff, ) if err != nil { t.Fatalf("reap expired runtime in winning transaction: %v", err) @@ -830,12 +782,7 @@ func TestConcurrentRuntimeLockReapersApplyRecoveryOnce(t *testing.T) { if contendErr != nil || contending != 0 { t.Fatalf("contending reaper count=%d error=%v, want 0 and nil", contending, contendErr) } - if err := fixture.Store.Execution().IntegrationCommitTxWithNotifications( - ctx, - reapTx, - txNotifications, - "commit winning reaper", - ); err != nil { + if err := reapTx.Commit(ctx, "commit winning reaper"); err != nil { t.Fatalf("commit winning reaper: %v", err) } var runtimeExists bool diff --git a/internal/storage/executionstore/agent_runtime_locks_store.go b/internal/storage/executionstore/agent_runtime_locks_store.go index a61ea4aa8..842cbdc0a 100644 --- a/internal/storage/executionstore/agent_runtime_locks_store.go +++ b/internal/storage/executionstore/agent_runtime_locks_store.go @@ -2,16 +2,13 @@ package executionstore import ( "context" - "encoding/json" "errors" "fmt" "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -36,40 +33,6 @@ type AgentRuntimeLockRenewal struct { LocalLeaseBudgetStartedAt time.Time } -func acquireAgentRuntimeLockTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, workerProcessID uuid.UUID, - leaseDuration time.Duration, -) (AgentRuntimeLockRecord, uuid.UUID, error) { - agent, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: projectID, ID: agentID}, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return AgentRuntimeLockRecord{}, uuid.Nil, storeerr.ErrAgentNotAdvanceable - } - return AgentRuntimeLockRecord{}, uuid.Nil, fmt.Errorf("lock agent for runtime acquisition: %w", err) - } - row, err := qtx.AcquireAgentRuntimeLock( - ctx, - dbsqlc.AcquireAgentRuntimeLockParams{ - ProjectID: projectID, - AgentID: agentID, - WorkerProcessID: workerProcessID, - LeaseDurationMicroseconds: leaseDuration.Microseconds(), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return AgentRuntimeLockRecord{}, uuid.Nil, storeerr.ErrAgentNotAdvanceable - } - if err != nil { - return AgentRuntimeLockRecord{}, uuid.Nil, fmt.Errorf("acquire agent runtime lock: %w", err) - } - return agentRuntimeLockRecordFromSQLC(row), agent.OrgID, nil -} - func validateAgentRuntimeLockLeaseDuration(leaseDuration time.Duration) error { if leaseDuration < MinimumAgentRuntimeLockLeaseDuration || leaseDuration > MaximumAgentRuntimeLockLeaseDuration { @@ -83,301 +46,60 @@ func validateAgentRuntimeLockLeaseDuration(leaseDuration time.Duration) error { } func (s *Store) EnsureRuntimeLockActive( - ctx context.Context, - projectID, agentID, runtimeID uuid.UUID, + ctx context.Context, projectID, agentID, runtimeID uuid.UUID, ) error { if projectID == uuid.Nil || agentID == uuid.Nil || runtimeID == uuid.Nil { return errors.New("project, agent, and runtime lock ids are required") } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return fmt.Errorf("begin runtime lock check: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - if err := ensureRuntimeLockActiveTx(ctx, tx, projectID, agentID, runtimeID); err != nil { - return err - } - if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("commit runtime lock check: %w", err) - } - return nil -} - -func lockAgentRuntimeForOwnedMutationTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeID uuid.UUID, -) error { - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: projectID, ID: agentID}, - ); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrRuntimeLockInactive - } - return fmt.Errorf("lock agent for runtime-owned mutation: %w", err) - } - if _, err := qtx.LockAgentRuntimeLockForOwnedMutation( - ctx, - dbsqlc.LockAgentRuntimeLockForOwnedMutationParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeID, - }, - ); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrRuntimeLockInactive - } - return fmt.Errorf("lock runtime for owned mutation: %w", err) - } - return nil -} - -func (s *Store) beginAgentRuntimeOwnedMutation( - ctx context.Context, - projectID, agentID, runtimeID uuid.UUID, -) (pgx.Tx, *dbsqlc.Queries, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return nil, nil, fmt.Errorf("begin runtime-owned mutation: %w", err) - } - qtx := dbsqlc.New(tx) - if err := lockAgentRuntimeForOwnedMutationTx(ctx, qtx, projectID, agentID, runtimeID); err != nil { - _ = tx.Rollback(ctx) - return nil, nil, err - } - return tx, qtx, nil -} - -func agentRuntimeLockActiveTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeID uuid.UUID, -) error { - active, err := qtx.AgentRuntimeLockIsActive( - ctx, - dbsqlc.AgentRuntimeLockIsActiveParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeID, - }, - ) - if err != nil { - return fmt.Errorf("check agent runtime lock: %w", err) - } - if !active { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if errors.Is(err, storeerr.ErrNotFound) { return storeerr.ErrRuntimeLockInactive } - return nil -} - -func ensureRuntimeLockActiveTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID, runtimeID uuid.UUID, -) error { - if projectID == uuid.Nil || agentID == uuid.Nil || runtimeID == uuid.Nil { - return errors.New("project, agent, and runtime lock ids are required") - } - qtx := dbsqlc.New(tx) - if err := lockAgentRuntimeForOwnedMutationTx(ctx, qtx, projectID, agentID, runtimeID); err != nil { - return err - } - return agentRuntimeLockActiveTx(ctx, qtx, projectID, agentID, runtimeID) -} - -func failQueuedRuntimeWorkTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeID uuid.UUID, - reason string, -) error { - if runtimeID == uuid.Nil { - return nil - } - params := dbsqlc.FailQueuedProcessesForRuntimeEndParams{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtimeID, - StateReasonCode: storeutil.TextFromEmpty(reason), - } - processRows, err := qtx.FailQueuedProcessesForRuntimeEnd(ctx, params) - if err != nil { - return fmt.Errorf("fail queued processes for ended runtime: %w", err) - } - for _, row := range processRows { - record := processRecordFromSQLC(row) - if err := completeProcessToolCallFromRecordTx( - ctx, - txNotifications, - tx, - qtx, - record, - nil, - reason, - ); err != nil { - return err - } - } - actionRows, err := qtx.FailQueuedProcessActionsForRuntimeEnd( - ctx, - dbsqlc.FailQueuedProcessActionsForRuntimeEndParams(params), - ) - if err != nil { - return fmt.Errorf("fail queued process actions for ended runtime: %w", err) - } - for _, row := range actionRows { - record := processActionRecordFromSQLC(row) - if err := completeTerminalProcessActionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - ProcessActionStateFailed, - reason, - ); err != nil { - return err - } - } - return nil -} - -const RuntimeToolInterruptedMessage = "Tool call was interrupted; its external outcome is unknown." - -func failRuntimeToolCallsTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - projectID, agentID, runtimeID uuid.UUID, - reason string, -) error { - result, err := marshalJSON(map[string]any{ - "error_code": reason, - "message": RuntimeToolInterruptedMessage, - }) - if err != nil { - return fmt.Errorf("marshal interrupted runtime tool result: %w", err) - } - contentParts, err := ToolResultContentParts(result) if err != nil { return err } - qtx := dbsqlc.New(tx) - rows, err := qtx.FailRuntimeToolCalls( - ctx, - dbsqlc.FailRuntimeToolCallsParams{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtimeID, - }, - ) - if err != nil { - return fmt.Errorf("fail interrupted runtime tool calls: %w", err) - } - for _, row := range rows { - record := toolCallRecordFromFailRuntimeSQLC(row) - record.ResultContentParts = contentParts - if _, err := appendToolResultEventTx(ctx, txNotifications, tx, record, nil); err != nil { - return err - } - } - if len(rows) == 0 { - return nil - } - metadata, err := marshalJSON(map[string]any{ - "reason": reason, - "runtime_lock_id": runtimeID, - }) - if err != nil { + defer func() { _ = unit.Rollback(ctx) }() + if err := h.FenceRuntime(ctx, runtimeID); err != nil { return err } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - Metadata: metadata, - }, - ); err != nil { - return fmt.Errorf("mark interrupted runtime tool wakeup: %w", err) - } - return nil + return unit.Commit(ctx, "runtime lock check") } +const RuntimeToolInterruptedMessage = "Tool call was interrupted; its external outcome is unknown." + func (s *Store) RenewAgentRuntimeLock( ctx context.Context, projectID, agentID, runtimeLockID uuid.UUID, leaseDuration time.Duration, ) (AgentRuntimeLockRenewal, error) { - if projectID == uuid.Nil || agentID == uuid.Nil || runtimeLockID == uuid.Nil { - return AgentRuntimeLockRenewal{}, errors.New("project, agent, and runtime lock ids are required") - } - if err := validateAgentRuntimeLockLeaseDuration(leaseDuration); err != nil { - return AgentRuntimeLockRenewal{}, err - } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { - return AgentRuntimeLockRenewal{}, fmt.Errorf("begin renew agent runtime lock: %w", err) + return AgentRuntimeLockRenewal{}, err } - defer func() { _ = tx.Rollback(ctx) }() - record, err := renewAgentRuntimeLockTx( + defer func() { _ = unit.Rollback(ctx) }() + renewal, err := unit.RenewRuntime( ctx, - dbsqlc.New(tx), - projectID, - agentID, + agentexecution.AgentRoute{CellID: unit.CellID(), ProjectID: projectID, AgentID: agentID}, runtimeLockID, leaseDuration, ) if err != nil { return AgentRuntimeLockRenewal{}, err } - if err := tx.Commit(ctx); err != nil { - return AgentRuntimeLockRenewal{}, fmt.Errorf("commit renew agent runtime lock: %w", err) - } - return record, nil -} - -func renewAgentRuntimeLockTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeLockID uuid.UUID, - leaseDuration time.Duration, -) (AgentRuntimeLockRenewal, error) { - if _, err := qtx.LockAgentRuntimeLockForRenewal( - ctx, - dbsqlc.LockAgentRuntimeLockForRenewalParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeLockID, - }, - ); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return AgentRuntimeLockRenewal{}, storeerr.ErrRuntimeLockInactive - } - return AgentRuntimeLockRenewal{}, fmt.Errorf("lock agent runtime renewal: %w", err) - } - localLeaseBudgetStartedAt := time.Now() //nolint:omnaralint // Local monotonic scheduling budget; never persisted. - row, err := qtx.RenewAgentRuntimeLock( - ctx, - dbsqlc.RenewAgentRuntimeLockParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeLockID, - LeaseDurationMicroseconds: leaseDuration.Microseconds(), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return AgentRuntimeLockRenewal{}, storeerr.ErrRuntimeLockInactive - } + row, err := dbsqlc.New(unit.DB()). + GetAgentRuntimeLockForRelease(ctx, + dbsqlc.GetAgentRuntimeLockForReleaseParams{ProjectID: projectID, + AgentID: agentID, + ID: runtimeLockID}) if err != nil { - return AgentRuntimeLockRenewal{}, fmt.Errorf("renew agent runtime lock: %w", err) + return AgentRuntimeLockRenewal{}, err + } + if err = unit.Commit(ctx, "renew runtime"); err != nil { + return AgentRuntimeLockRenewal{}, err } return AgentRuntimeLockRenewal{ RuntimeLock: agentRuntimeLockRecordFromSQLC(row), - LocalLeaseBudgetStartedAt: localLeaseBudgetStartedAt, + LocalLeaseBudgetStartedAt: renewal.LocalStartedAt, }, nil } @@ -385,80 +107,16 @@ func (s *Store) ReleaseAgentRuntimeLock( ctx context.Context, projectID, agentID, runtimeLockID uuid.UUID, ) error { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return fmt.Errorf("begin release agent runtime lock: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if err := lockAgentWithParentTx(ctx, tx, qtx, projectID, agentID); err != nil { - if errors.Is(err, storeerr.ErrNotFound) { - return storeerr.ErrRuntimeLockInactive - } - return fmt.Errorf("lock agent for runtime lock release: %w", err) - } - releaseParams := dbsqlc.GetAgentRuntimeLockForReleaseParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeLockID, - } - if _, err := qtx.GetAgentRuntimeLockForRelease(ctx, releaseParams); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrRuntimeLockInactive - } - return fmt.Errorf("get agent runtime lock for release: %w", err) - } - if err := recoverRuntimeModelCallContextTx( - ctx, - txNotifications, - tx, - qtx, - projectID, - agentID, - runtimeLockID, - runtimeModelCallRecoveryEvidence{ - Code: "runtime_released_before_model_result_acceptance", - Message: "runtime released before the model result was durably accepted", - }, - s.modelCallRetryDelay, - ); err != nil { - return err - } - if err := failRuntimeToolCallsTx( - ctx, - txNotifications, - tx, - projectID, - agentID, - runtimeLockID, - "runtime_lock_released", - ); err != nil { - return err - } - changed, err := qtx.ReleaseAgentRuntimeLock( - ctx, - dbsqlc.ReleaseAgentRuntimeLockParams{ - ProjectID: projectID, - AgentID: agentID, - ID: runtimeLockID, - }, - ) - if err != nil { - return fmt.Errorf("release agent runtime lock: %w", err) - } - if changed != 1 { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if errors.Is(err, storeerr.ErrNotFound) { return storeerr.ErrRuntimeLockInactive } - if err := qtx.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - Metadata: json.RawMessage(`{"reason":"runtime_release"}`), - }); err != nil { - return fmt.Errorf("reconcile agent wakeup after runtime release: %w", err) + if err != nil { + return err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "release agent runtime lock"); err != nil { + defer func() { _ = unit.Rollback(ctx) }() + if err = h.Release(ctx, runtimeLockID); err != nil { return err } - return nil + return unit.Commit(ctx, "release runtime") } diff --git a/internal/storage/executionstore/agent_runtime_recovery_store.go b/internal/storage/executionstore/agent_runtime_recovery_store.go deleted file mode 100644 index af98f8821..000000000 --- a/internal/storage/executionstore/agent_runtime_recovery_store.go +++ /dev/null @@ -1,159 +0,0 @@ -package executionstore - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "time" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/modelprotocol" - "github.com/omnara-ai/omnara/internal/notifications" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/storeerr" -) - -type runtimeModelCallRecoveryEvidence struct { - Code string - Message string -} - -func recoverRuntimeModelCallContextTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeLockID uuid.UUID, - evidence runtimeModelCallRecoveryEvidence, - retryBackoff func(int, string) time.Duration, -) error { - contextID, err := qtx.GetLiveModelCallContextForRuntime( - ctx, - dbsqlc.GetLiveModelCallContextForRuntimeParams{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtimeLockID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return nil - } - if err != nil { - return fmt.Errorf("load live model call context for runtime recovery: %w", err) - } - contextRecord, err := loadModelCallContextByID(ctx, qtx, projectID, agentID, contextID) - if err != nil { - return fmt.Errorf("load live model call context for runtime recovery: %w", err) - } - details, err := marshalJSON(map[string]any{ - "attempt_number": contextRecord.AttemptNumber, - "code": evidence.Code, - "message": evidence.Message, - "outcome_ambiguous": true, - "runtime_lock_id": runtimeLockID, - }) - if err != nil { - return fmt.Errorf("marshal runtime recovery error: %w", err) - } - if contextRecord.AttemptNumber > MaxModelCallRetriesPerOperation { - return terminalizeExhaustedRuntimeModelCallContextTx( - ctx, - txNotifications, - tx, - qtx, - projectID, - agentID, - runtimeLockID, - contextRecord, - evidence, - details, - ) - } - changed, err := qtx.InterruptRuntimeModelCallContextForRetry( - ctx, - dbsqlc.InterruptRuntimeModelCallContextForRetryParams{ - ProjectID: projectID, - AgentID: agentID, - ID: contextRecord.ID, - RuntimeLockID: runtimeLockID, - ErrorKind: string(modelprotocol.ErrorKindRuntime), - ErrorCode: evidence.Code, - ErrorMessage: evidence.Message, - ErrorDetails: details, - RetryDelayMicroseconds: retryBackoff( - contextRecord.AttemptNumber, - contextRecord.ID.String(), - ).Microseconds(), - }, - ) - if err != nil { - return fmt.Errorf("interrupt runtime model call context: %w", err) - } - if changed != 1 { - return storeerr.ErrStateTransitionConflict - } - return nil -} - -func terminalizeExhaustedRuntimeModelCallContextTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeLockID uuid.UUID, - contextRecord ModelCallContextRecord, - evidence runtimeModelCallRecoveryEvidence, - details json.RawMessage, -) error { - if contextRecord.ID == uuid.Nil { - return errors.New("model call context is required for runtime recovery") - } - switch contextRecord.OperationKind { - case ModelCallOperationNormal: - if _, err := recordTerminalModelCallFailureTx( - ctx, - txNotifications, - tx, - qtx, - RecordModelCallErrorAndCompleteContextInput{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtimeLockID, - ModelCallContextID: contextRecord.ID, - ErrorKind: modelprotocol.ErrorKindRuntime, - ErrorCode: evidence.Code, - ErrorMessage: evidence.Message, - ErrorDetails: details, - }, - modelCallContextRuntimeTeardown, - ModelCallOperationNormal, - ); err != nil { - return fmt.Errorf("terminalize exhausted normal model call: %w", err) - } - case ModelCallOperationCompaction: - if err := recordTerminalCompactionFailureTx( - ctx, - txNotifications, - tx, - qtx, - RecordTerminalCompactionFailureInput{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtimeLockID, - ModelCallContextID: contextRecord.ID, - ErrorKind: modelprotocol.ErrorKindRuntime, - ErrorCode: evidence.Code, - ErrorMessage: evidence.Message, - ErrorDetails: details, - }, - modelCallContextRuntimeTeardown, - ); err != nil { - return fmt.Errorf("terminalize exhausted compaction model call: %w", err) - } - default: - return fmt.Errorf("unsupported model call operation %q", contextRecord.OperationKind) - } - return nil -} diff --git a/internal/storage/executionstore/agent_runtime_store.go b/internal/storage/executionstore/agent_runtime_store.go index 0ab4638ea..7c422fe04 100644 --- a/internal/storage/executionstore/agent_runtime_store.go +++ b/internal/storage/executionstore/agent_runtime_store.go @@ -9,8 +9,8 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/integrationdefinition" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -78,69 +78,49 @@ type IntegrationTargetDisplay struct { DisplayName string `json:"display_name,omitempty"` } -func insertAdmittedAgentTx( - ctx context.Context, - tx pgx.Tx, - qtx *dbsqlc.Queries, - input insertAgentInput, -) (AgentRecord, bool, error) { - row, err := qtx.InsertAgent(ctx, dbsqlc.InsertAgentParams{ - ID: storeutil.IDFromNil(input.ID), - OrgID: input.OrgID, - ProjectID: input.ProjectID, - AgentProfileID: storeutil.IDFromNil(input.AgentProfileID), - Name: input.Name, - CurrentConfigID: input.CurrentConfigID, - IdempotencyKey: storeutil.TextFromEmpty(input.IdempotencyKey), - ParentAgentID: storeutil.IDFromNil(input.ParentAgentID), - SubagentKey: input.SubagentKey, - ArchiveAfterIdleMinutes: storeutil.Int32Ptr(input.ArchiveAfterIdleMinutes), - }) - if err == nil { - record := agentRecordFromInsertSQLC(row) - record.Created = true - if err := lockResourceCreation(ctx, qtx, resourceAgents, input.ProjectID.String()); err != nil { - return AgentRecord{}, false, err +func insertAdmittedAgentTx(ctx context.Context, unit *agentexecution.Unit, qtx *dbsqlc.Queries, + input insertAgentInput) (AgentRecord, bool, error) { + if input.IdempotencyKey != "" { + record, err := loadAgentByIdempotencyKeyTx(ctx, unit.DB(), input.ProjectID, input.IdempotencyKey) + if err == nil { + return record, false, nil } - limits, err := resolveResourceLimits(ctx, qtx, input.OrgID) - if err != nil { + if !errors.Is(err, pgx.ErrNoRows) { return AgentRecord{}, false, err } - agentCount, err := qtx.CountActiveAgentsForProject( - ctx, - dbsqlc.CountActiveAgentsForProjectParams{ProjectID: input.ProjectID}, - ) - if err != nil { - return AgentRecord{}, false, fmt.Errorf("count active agents: %w", err) - } - if agentCount > limits.MaxActiveAgentsPerProject { - return AgentRecord{}, false, resourceLimitExceeded( - "active agents", - limits.MaxActiveAgentsPerProject, - ) - } - return record, true, nil } - if !errors.Is(err, pgx.ErrNoRows) { - if storeutil.IsUniqueViolation(err) { - return AgentRecord{}, false, storeerr.ErrIdempotencyConflict - } - return AgentRecord{}, false, fmt.Errorf("create agent: %w", err) + if input.ID == uuid.Nil { + input.ID = uuid.New() } - if input.IdempotencyKey == "" { - return AgentRecord{}, false, fmt.Errorf("create agent: %w", storeerr.ErrIdempotencyConflict) + _, err := unit.CreateAgent(ctx, agentexecution.CreateAgentInput{ID: input.ID, ProjectID: input.ProjectID, + ParentID: input.ParentAgentID, ConfigID: input.CurrentConfigID, ProfileID: input.AgentProfileID, + Name: input.Name, SubagentKey: input.SubagentKey, IdempotencyKey: input.IdempotencyKey, + ArchiveAfterIdleMinutes: storeutil.Int32Ptr(input.ArchiveAfterIdleMinutes)}) + if err != nil { + return AgentRecord{}, false, err } - - record, err := loadAgentByIdempotencyKeyTx(ctx, tx, input.ProjectID, input.IdempotencyKey) + limits, err := resolveResourceLimits(ctx, qtx, input.OrgID) if err != nil { return AgentRecord{}, false, err } - return record, false, nil + count, err := qtx.CountActiveAgentsForProject( + ctx, + dbsqlc.CountActiveAgentsForProjectParams{ProjectID: input.ProjectID}, + ) + if err != nil { + return AgentRecord{}, false, err + } + if count > limits.MaxActiveAgentsPerProject { + return AgentRecord{}, false, resourceLimitExceeded("active agents", limits.MaxActiveAgentsPerProject) + } + record, err := loadAgentInProjectTx(ctx, unit.DB(), input.ProjectID, input.ID) + record.Created = true + return record, err == nil, err } func loadAgentByIdempotencyKeyTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID uuid.UUID, idempotencyKey string, ) (AgentRecord, error) { @@ -155,7 +135,7 @@ func loadAgentByIdempotencyKeyTx( return agentRecordFromIdempotencySQLC(row), nil } -func loadAgentInProjectTx(ctx context.Context, tx pgx.Tx, projectID, id uuid.UUID) (AgentRecord, error) { +func loadAgentInProjectTx(ctx context.Context, tx dbsqlc.DBTX, projectID, id uuid.UUID) (AgentRecord, error) { row, err := dbsqlc.New(tx). GetAgentInProject(ctx, dbsqlc.GetAgentInProjectParams{ProjectID: projectID, ID: id}) if err != nil { @@ -287,7 +267,10 @@ func (s *Store) attachAgentActivity( func (s *Store) GetAgentActivity(ctx context.Context, projectID, agentID uuid.UUID) (AgentActivity, error) { rows, err := s.q.ListAgentActivityForAgents( ctx, - dbsqlc.ListAgentActivityForAgentsParams{ProjectIds: []uuid.UUID{projectID}, AgentIds: []uuid.UUID{agentID}}, + dbsqlc.ListAgentActivityForAgentsParams{ + ProjectIds: []uuid.UUID{projectID}, + AgentIds: []uuid.UUID{agentID}, + }, ) if err != nil { return AgentActivity{}, fmt.Errorf("get agent activity: %w", err) @@ -301,7 +284,12 @@ func (s *Store) GetAgentActivity(ctx context.Context, projectID, agentID uuid.UU func agentActivityFromRow(row dbsqlc.ListAgentActivityForAgentsRow) AgentActivity { return AgentActivity{ State: agentActivityState( - row.State == string(AgentStateArchived), row.HasOpenQuestion, row.HasOpenPermission, row.IsRunning, + row.State == string( + AgentStateArchived, + ), + row.HasOpenQuestion, + row.HasOpenPermission, + row.IsRunning, ), LastActivityAt: row.LastActivityAt, } @@ -433,17 +421,18 @@ func (s *Store) archiveAgentOnce( projectID, agentID uuid.UUID, actor *ActorParams, ) (AgentRecord, []MachineRecord, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentRecord{}, nil, fmt.Errorf("begin archive agent: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) if err := enterActiveAgentProjectTx(ctx, tx, qtx, projectID); err != nil { return AgentRecord{}, nil, err } - machines, err := archiveAgentTreeTx(ctx, tx, qtx, txNotifications, projectID, agentID, actor) + machines, err := archiveAgentTreeTx(ctx, unit, qtx, projectID, agentID, actor) if err != nil { return AgentRecord{}, nil, err } @@ -451,12 +440,7 @@ func (s *Store) archiveAgentOnce( if err != nil { return AgentRecord{}, nil, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - fmt.Sprintf("archive agent %s", agentID.String()), - ); err != nil { + if err := unit.Commit(ctx, fmt.Sprintf("archive agent %s", agentID.String())); err != nil { return AgentRecord{}, nil, err } return agent, machines, nil @@ -464,9 +448,8 @@ func (s *Store) archiveAgentOnce( func archiveAgentTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, - txNotifications *notifications.TxNotifications, projectID, agentID uuid.UUID, actor *ActorParams, ) ([]MachineRecord, error) { @@ -474,69 +457,22 @@ func archiveAgentTx( if err != nil { return nil, err } - rows, err := qtx.ArchiveAgent(ctx, dbsqlc.ArchiveAgentParams{ProjectID: projectID, ID: agentID}) + h, err := unit.Handle(projectID, agentID) if err != nil { - return nil, fmt.Errorf("archive agent: %w", err) - } - if rows == 0 { - return nil, storeerr.ErrNotFound - } - - if err := qtx.DeleteAgentIntegrationSubscriptions(ctx, dbsqlc.DeleteAgentIntegrationSubscriptionsParams{ - ProjectID: projectID, AgentID: agentID, - }); err != nil { - return nil, fmt.Errorf("remove archived agent subscriptions: %w", err) - } - if _, err := qtx.CancelQueuedBacklogInputsForAgent(ctx, dbsqlc.CancelQueuedBacklogInputsForAgentParams{ - ProjectID: projectID, - AgentID: agentID, - }); err != nil { - return nil, fmt.Errorf("cancel archived agent queued backlog inputs: %w", err) - } - if _, err := qtx.DeleteCronTriggersForAgent(ctx, dbsqlc.DeleteCronTriggersForAgentParams{ - ProjectID: projectID, - AgentID: &agentID, - }); err != nil { - return nil, fmt.Errorf("delete archived agent cron triggers: %w", err) + return nil, err } - if err := completeExecutionRevokedProcessesTx( - ctx, - txNotifications, - tx, - qtx, - executionRevokedProcessScope{projectID: projectID, agentID: agentID}, - "agent_archived", - ); err != nil { + if _, err = h.Archive(ctx, actorID); err != nil { return nil, err } - if _, err := cancelAgentTx( + machineRows, err := qtx.MarkArchivedAgentPoolMachinesDeleting( ctx, - txNotifications, - tx, - qtx, - cancelAgentTxInput{ - ProjectID: projectID, - AgentID: agentID, - ActorID: actorID, - ReasonCode: "agent_archived", - ModelCallMessage: "The model call was stopped because the agent was archived.", - CancelRuntimeWithoutContinuableTurn: true, + dbsqlc.MarkArchivedAgentPoolMachinesDeletingParams{ + ProjectID: projectID, + AgentID: agentID, + LifecycleReasonCode: storeutil.TextFromEmpty("agent_archived_cleanup"), + LifecycleReasonMessage: "cleaning up machine after agent archived", }, - ); err != nil { - return nil, err - } - if _, err := qtx.DeleteAgentWakeup(ctx, dbsqlc.DeleteAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - }); err != nil { - return nil, fmt.Errorf("delete archived agent wakeup: %w", err) - } - machineRows, err := qtx.MarkArchivedAgentPoolMachinesDeleting(ctx, dbsqlc.MarkArchivedAgentPoolMachinesDeletingParams{ - ProjectID: projectID, - AgentID: agentID, - LifecycleReasonCode: storeutil.TextFromEmpty("agent_archived_cleanup"), - LifecycleReasonMessage: "cleaning up machine after agent archived", - }) + ) if err != nil { return nil, fmt.Errorf("mark archived agent pool machines deleting: %w", err) } @@ -554,7 +490,11 @@ func archiveAgentTx( return machines, nil } -func listActiveAgentTreeTx(ctx context.Context, qtx *dbsqlc.Queries, projectID, rootID uuid.UUID) ([]uuid.UUID, error) { +func listActiveAgentTreeTx( + ctx context.Context, + qtx *dbsqlc.Queries, + projectID, rootID uuid.UUID, +) ([]uuid.UUID, error) { tree := []uuid.UUID{rootID} for index := 0; index < len(tree); index++ { parentID := tree[index] @@ -588,11 +528,12 @@ func sameAgentIDSet(left, right []uuid.UUID) bool { func lockAgentTreeForArchiveTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, root AgentRecord, tree []uuid.UUID, ) error { + tx := unit.DB() if err := lifecyclelock.AgentSources(ctx, tx, tree...); err != nil { return err } @@ -600,13 +541,19 @@ func lockAgentTreeForArchiveTx( for _, agentID := range tree { machineIDs, err := qtx.ListAttachedAgentPoolMachineIDsForLifecycle( ctx, - dbsqlc.ListAttachedAgentPoolMachineIDsForLifecycleParams{ProjectID: root.ProjectID, AgentID: agentID}, + dbsqlc.ListAttachedAgentPoolMachineIDsForLifecycleParams{ + ProjectID: root.ProjectID, + AgentID: agentID, + }, ) if err != nil { return fmt.Errorf("list archived agent pool machines: %w", err) } for _, machineID := range machineIDs { - machineRefs = append(machineRefs, lifecyclelock.MachineRef{OrgID: root.OrgID, MachineID: machineID}) + machineRefs = append( + machineRefs, + lifecyclelock.MachineRef{OrgID: root.OrgID, MachineID: machineID}, + ) } } if err := lifecyclelock.Machines(ctx, tx, machineRefs); err != nil { @@ -617,9 +564,12 @@ func lockAgentTreeForArchiveTx( agentRefs = append(agentRefs, lifecyclelock.AgentRef{ProjectID: root.ProjectID, AgentID: agentID}) } if root.ParentAgentID != uuid.Nil { - agentRefs = append(agentRefs, lifecyclelock.AgentRef{ProjectID: root.ProjectID, AgentID: root.ParentAgentID}) + agentRefs = append( + agentRefs, + lifecyclelock.AgentRef{ProjectID: root.ProjectID, AgentID: root.ParentAgentID}, + ) } - if err := lifecyclelock.Agents(ctx, tx, agentRefs); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { return err } lockedTree, err := listActiveAgentTreeTx(ctx, qtx, root.ProjectID, root.ID) @@ -640,7 +590,12 @@ type lockedAgentTree struct { // enterActiveAgentProjectTx takes the organization and project lifecycle // gates before an agent tree is locked, so archival and stop paths order // against project and organization deletion the same way ArchiveAgent does. -func enterActiveAgentProjectTx(ctx context.Context, tx pgx.Tx, qtx *dbsqlc.Queries, projectID uuid.UUID) error { +func enterActiveAgentProjectTx( + ctx context.Context, + tx dbsqlc.DBTX, + qtx *dbsqlc.Queries, + projectID uuid.UUID, +) error { project, err := loadProjectTx(ctx, qtx, projectID) if err != nil { return err @@ -650,10 +605,11 @@ func enterActiveAgentProjectTx(ctx context.Context, tx pgx.Tx, qtx *dbsqlc.Queri func lockAgentTreeTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, projectID, agentID uuid.UUID, ) (lockedAgentTree, error) { + tx := unit.DB() root, err := loadAgentInProjectTx(ctx, tx, projectID, agentID) if err != nil { return lockedAgentTree{}, err @@ -662,7 +618,7 @@ func lockAgentTreeTx( if err != nil { return lockedAgentTree{}, err } - if err := lockAgentTreeForArchiveTx(ctx, tx, qtx, root, tree); err != nil { + if err := lockAgentTreeForArchiveTx(ctx, unit, qtx, root, tree); err != nil { return lockedAgentTree{}, err } root, err = loadAgentInProjectTx(ctx, tx, projectID, agentID) @@ -674,15 +630,21 @@ func lockAgentTreeTx( func archiveLockedAgentTreeTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, - txNotifications *notifications.TxNotifications, locked lockedAgentTree, actor *ActorParams, ) ([]MachineRecord, error) { var machines []MachineRecord for index := len(locked.ids) - 1; index >= 0; index-- { - released, err := archiveAgentTx(ctx, tx, qtx, txNotifications, locked.root.ProjectID, locked.ids[index], actor) + released, err := archiveAgentTx( + ctx, + unit, + qtx, + locked.root.ProjectID, + locked.ids[index], + actor, + ) if err != nil { return nil, err } @@ -693,15 +655,15 @@ func archiveLockedAgentTreeTx( func archiveAgentTreeTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, - txNotifications *notifications.TxNotifications, projectID, agentID uuid.UUID, actor *ActorParams, ) ([]MachineRecord, error) { - locked, err := lockAgentTreeTx(ctx, tx, qtx, projectID, agentID) + + locked, err := lockAgentTreeTx(ctx, unit, qtx, projectID, agentID) if err != nil { return nil, err } - return archiveLockedAgentTreeTx(ctx, tx, qtx, txNotifications, locked, actor) + return archiveLockedAgentTreeTx(ctx, unit, qtx, locked, actor) } diff --git a/internal/storage/executionstore/agent_turns_store.go b/internal/storage/executionstore/agent_turns_store.go index fd714dc02..7350803ba 100644 --- a/internal/storage/executionstore/agent_turns_store.go +++ b/internal/storage/executionstore/agent_turns_store.go @@ -9,10 +9,7 @@ import ( "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -91,141 +88,6 @@ const ( maxAgentTurnsReadLimit = MaxAgentTurnsReadPageLimit + 1 ) -func updateAgentTurnLatestEventTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID, turnID, latestEventID, latestSemanticEventID uuid.UUID, -) error { - if projectID == uuid.Nil || agentID == uuid.Nil || turnID == uuid.Nil || latestEventID == uuid.Nil { - return errors.New("project id, agent id, turn id, and latest event id are required") - } - return updateAgentTurnLatestEventQuery( - ctx, - dbsqlc.New(tx), - projectID, - agentID, - turnID, - latestEventID, - latestSemanticEventID, - ) -} - -func updateAgentTurnLatestEventQuery( - ctx context.Context, - queries *dbsqlc.Queries, - projectID, agentID, turnID, latestEventID, latestSemanticEventID uuid.UUID, -) error { - changed, err := queries.UpdateAgentTurnLatestEvent( - ctx, - dbsqlc.UpdateAgentTurnLatestEventParams{ - ProjectID: projectID, - AgentID: agentID, - ID: turnID, - LatestEventID: latestEventID, - LatestSemanticEventID: storeutil.IDFromNil(latestSemanticEventID), - }, - ) - if err != nil { - return fmt.Errorf("update agent turn latest event: %w", err) - } - if changed != 1 { - return storeerr.ErrStateTransitionConflict - } - return nil -} - -func createSingleEventAgentTurnTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, turnID uuid.UUID, - event events.Event, -) (AgentTurnRecord, error) { - sequence, err := qtx.NextTurnSequence(ctx, dbsqlc.NextTurnSequenceParams{ProjectID: projectID, AgentID: agentID}) - if err != nil { - return AgentTurnRecord{}, fmt.Errorf("next turn sequence: %w", err) - } - row, err := qtx.InsertAgentTurn( - ctx, - dbsqlc.InsertAgentTurnParams{ - ID: turnID, - ProjectID: projectID, - AgentID: agentID, - TurnSequence: sequence, - LatestEventID: event.ID, - LatestSemanticEventID: event.ID, - }, - ) - if err != nil { - return AgentTurnRecord{}, fmt.Errorf("insert agent turn: %w", err) - } - return agentTurnRecordFromInsertSQLC(row), nil -} - -func appendEventToCurrentOrNewAgentTurnTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - eventInput AppendTypedAgentEventInput, - semantic bool, -) (TypedAgentEventRecord, AgentTurnRecord, bool, error) { - projectID := eventInput.ProjectID - agentID := eventInput.AgentID - turn, err := qtx.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: projectID, AgentID: agentID}, - ) - if err == nil { - eventInput.TurnID = turn.ID - eventInput.IsOpeningEvent = false - eventRecord, err := appendTypedAgentEventTx(ctx, txNotifications, tx, eventInput) - if err != nil { - return TypedAgentEventRecord{}, AgentTurnRecord{}, false, err - } - semanticEventID := uuid.Nil - if semantic { - semanticEventID = eventRecord.Event.ID - } - if err := updateAgentTurnLatestEventQuery( - ctx, - qtx, - projectID, - agentID, - turn.ID, - eventRecord.Event.ID, - semanticEventID, - ); err != nil { - return TypedAgentEventRecord{}, AgentTurnRecord{}, false, err - } - return eventRecord, agentTurnRecordFromCurrentContinuableSQLC(turn), false, nil - } - if !errors.Is(err, pgx.ErrNoRows) { - return TypedAgentEventRecord{}, AgentTurnRecord{}, false, fmt.Errorf("load current continuable turn: %w", err) - } - turnID, err := uuid.NewV7() - if err != nil { - return TypedAgentEventRecord{}, AgentTurnRecord{}, false, fmt.Errorf("generate turn id: %w", err) - } - eventInput.TurnID = turnID - eventInput.IsOpeningEvent = true - eventRecord, err := appendTypedAgentEventTx(ctx, txNotifications, tx, eventInput) - if err != nil { - return TypedAgentEventRecord{}, AgentTurnRecord{}, false, err - } - turnRecord, err := createSingleEventAgentTurnTx( - ctx, - qtx, - projectID, - agentID, - turnID, - eventRecord.Event, - ) - if err != nil { - return TypedAgentEventRecord{}, AgentTurnRecord{}, false, err - } - return eventRecord, turnRecord, true, nil -} - func (s *Store) ListAgentEventsForRead( ctx context.Context, projectID, agentID uuid.UUID, diff --git a/internal/storage/executionstore/agent_work_integration_test.go b/internal/storage/executionstore/agent_work_integration_test.go new file mode 100644 index 000000000..7cffb4242 --- /dev/null +++ b/internal/storage/executionstore/agent_work_integration_test.go @@ -0,0 +1,511 @@ +//go:build integration + +package executionstore_test + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/interactionform" + "github.com/omnara-ai/omnara/internal/processcmd" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + "github.com/omnara-ai/omnara/internal/toolpermission" + "github.com/stretchr/testify/require" +) + +func ownedAdvanceInput(f processDaemonFixture, allowModel bool) executionstore.AdvanceOwnedAgentWorkInput { + return executionstore.AdvanceOwnedAgentWorkInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, AllowModelWork: allowModel, + } +} + +func ownedInput( + t *testing.T, + ctx context.Context, + f processDaemonFixture, + mode executionstore.AgentInputDeliveryMode, + key string, +) executionstore.AgentInputRecord { + t.Helper() + input, _, _, err := f.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, AgentID: f.AgentID, Actor: mustOmnaraActorParams(t, f.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"next input"}]`), + DeliveryMode: mode, IdempotencyKey: key, + }) + require.NoError(t, err) + return input +} + +func ownedCompleteTool(t *testing.T, ctx context.Context, f processDaemonFixture, id uuid.UUID) { + t.Helper() + _, err := f.Store.Execution().CompleteToolCall(ctx, executionstore.CompleteToolCallInput{ + ProjectID: testProjectID, AgentID: f.AgentID, ID: id, RuntimeLockID: f.Lock.ID, + Outcome: executionstore.ToolResultOutcomeSucceeded, + ResultContentParts: json.RawMessage(`[{"type":"text","text":"done"}]`), + }) + require.NoError(t, err) +} + +func ownedModelClaim( + t *testing.T, + ctx context.Context, + f processDaemonFixture, + work executionstore.ClaimedAgentWork, +) executionstore.ModelCallClaim { + t.Helper() + _, err := f.Store.Execution().CaptureAgentConfigForModelContext(ctx, testProjectID, f.AgentID) + require.NoError(t, err) + prepared1, err := f.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, + OpeningInputIDs: work.Model.InputIDs, + SourceModelCallContextID: work.Model.SourceModelCallContextID, + SourceModelOutputID: work.Model.SourceModelOutputID, + }) + claim := prepared1.Claim + + require.NoError(t, err) + require.True(t, claim.Claimed) + return claim +} + +func requireOwnedLease(t *testing.T, ctx context.Context, f processDaemonFixture, exists bool) { + t.Helper() + var found bool + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT EXISTS (SELECT 1 FROM agent_runtime_locks WHERE agent_id=$1 AND id=$2)`, f.AgentID, f.Lock.ID, + ).Scan(&found)) + require.Equal(t, exists, found) +} + +func TestAdvanceOwnedAgentWorkPrecedenceMatchesGlobalClaim(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name string + toolType string + blocked, complete, steering, queued, ended bool + kind executionstore.AgentWorkKind + modelKind executionstore.ModelWorkKind + }{ + {name: "runnable tools before steering", steering: true, queued: true, kind: executionstore.AgentWorkTool}, + {name: "permission blocks steering", blocked: true, steering: true}, + {name: "custom tool blocks steering", toolType: "custom", steering: true}, + { + name: "steering before model continuation", + complete: true, + steering: true, + queued: true, + kind: executionstore.AgentWorkModel, + modelKind: executionstore.ModelWorkStart, + }, + { + name: "model continuation before queued input", + complete: true, + queued: true, + kind: executionstore.AgentWorkModel, + modelKind: executionstore.ModelWorkContinue, + }, + { + name: "completed tools continue", + complete: true, + kind: executionstore.AgentWorkModel, + modelKind: executionstore.ModelWorkContinue, + }, + { + name: "queued input after completed turn", + complete: true, + ended: true, + queued: true, + kind: executionstore.AgentWorkModel, + modelKind: executionstore.ModelWorkStart, + }, + {name: "idle after completed turn", complete: true, ended: true}, + } { + for _, owned := range []bool{false, true} { + name := test.name + "/global" + if owned { + name = test.name + "/owned" + } + t.Run(name, func(t *testing.T) { + t.Parallel() + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_precedence") + + item := builtInProcessToolCallBatchItem("first", "read_process") + if test.toolType != "" { + item.ToolType = test.toolType + } + item.Allowed = !test.blocked + toolID := createToolCallBatchForProcessTest(t, + ctx, + f, + "owned_precedence", + []processToolCallBatchItem{item})[0] + contextID := modelContextIDForProcessToolCallTest(t, ctx, f, toolID) + outputID := modelOutputIDForToolCall(t, ctx, f.Store, f.AgentID, toolID) + if test.blocked { + createPermissionInteractionForTest( + t, + ctx, + f, + toolID, + permissionRequestForStorageTest(t, "read_process"), + ) + } + if test.complete { + ownedCompleteTool(t, ctx, f, toolID) + } + if test.ended { + work, continued, err := f.Store.Execution(). + AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + claim := ownedModelClaim(t, ctx, f, work) + createModelOutputEventForTurnTest( + t, ctx, f, work.Model.TurnID, claim.Context.ID, "owned_end", "end_turn", "", f.Now, + ) + } + var steering, queued executionstore.AgentInputRecord + if test.steering { + steering = ownedInput(t, ctx, f, executionstore.DeliveryModeSteering, "steering") + } + if test.queued { + queued = ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "queued") + } + var work executionstore.ClaimedAgentWork + var err error + if owned { + var continued bool + work, continued, err = f.Store.Execution(). + AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.Equal(t, test.kind != executionstore.AgentWorkNone, continued) + requireOwnedLease(t, ctx, f, continued) + if continued { + require.Equal(t, f.Lock.ID, work.RuntimeLock.ID) + } + } else { + require.NoError(t, f.Store.Execution().ReleaseAgentRuntimeLock(ctx, testProjectID, f.AgentID, f.Lock.ID)) + work, _, err = f.Store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + if test.kind == executionstore.AgentWorkNone { + require.ErrorIs(t, err, storeerr.ErrNoClaimableAgentWakeup) + } else { + require.NoError(t, err) + } + } + require.Equal(t, test.kind, work.Kind) + if work.Kind != executionstore.AgentWorkNone { + require.Zero(t, countAgentWakeups(t, ctx, f.Store, f.AgentID)) + } + if work.Kind == executionstore.AgentWorkModel { + require.Equal(t, test.modelKind, work.Model.Kind) + if test.modelKind == executionstore.ModelWorkContinue { + require.Equal(t, contextID, work.Model.SourceModelCallContextID) + require.Equal(t, outputID, work.Model.SourceModelOutputID) + } else if test.steering { + require.Equal(t, []uuid.UUID{steering.ID}, work.Model.InputIDs) + } else { + require.Equal(t, []uuid.UUID{queued.ID}, work.Model.InputIDs) + } + } + if test.queued && !test.ended { + var state string + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT state FROM agent_inputs WHERE id=$1`, queued.ID, + ).Scan(&state)) + require.Equal(t, "received", state) + } + }) + } + } +} + +func TestAdvanceOwnedAgentWorkRunnableSiblingAndModelBudget(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_siblings") + + blocked := builtInProcessToolCallBatchItem("blocked", "read_process") + blocked.Allowed = false + ids := createToolCallBatchForProcessTest(t, ctx, f, "siblings", []processToolCallBatchItem{ + builtInProcessToolCallBatchItem("ready", "read_process"), blocked, + }) + interaction := createPermissionInteractionForTest( + t, ctx, f, ids[1], permissionRequestForStorageTest(t, "read_process"), + ) + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, false)) + require.NoError(t, err) + require.True(t, continued) + require.Equal(t, executionstore.AgentWorkTool, work.Kind) + ownedCompleteTool(t, ctx, f, ids[0]) + _, err = f.Store.Execution().ResolveAgentInteraction(ctx, executionstore.ResolveAgentInteractionInput{ + ProjectID: testProjectID, AgentID: f.AgentID, ID: interaction.ID, + Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{ + OptionIndices: []int{toolpermission.AllowOptionIndex}, + }}}, + Actor: mustOmnaraActorParams(t, f.UserID), + }) + require.NoError(t, err) + ownedCompleteTool(t, ctx, f, ids[1]) + steering := ownedInput(t, ctx, f, executionstore.DeliveryModeSteering, "budget-steering") + _, continued, err = f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, false)) + require.NoError(t, err) + require.False(t, continued) + requireOwnedLease(t, ctx, f, false) + var state string + require.NoError( + t, + f.Store.pool.QueryRow(ctx, `SELECT state FROM agent_inputs WHERE id=$1`, steering.ID).Scan(&state), + ) + require.Equal(t, "received", state) + work, found, err := f.Store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, []uuid.UUID{steering.ID}, work.Model.InputIDs) +} + +func TestAdvanceOwnedAgentWorkRetryAndUnfinishedRecovery(t *testing.T) { + for _, unfinished := range []bool{false, true} { + name := "future retry" + if unfinished { + name = "unfinished model" + } + t.Run(name, func(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_retry") + + ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "initial") + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + claim := ownedModelClaim(t, ctx, f, work) + if !unfinished { + _, err := f.Store.Execution().RecordRetryableModelCallFailure( + ctx, executionstore.RecordRecoverableModelCallFailureInput{ + ProjectID: testProjectID, + AgentID: f.AgentID, + RuntimeLockID: f.Lock.ID, + ModelCallContextID: claim.Context.ID, + ErrorKind: "transient", ErrorMessage: "retry later", RetryDelay: time.Hour, + }, + ) + require.NoError(t, err) + } + ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "waiting") + _, continued, err = f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.False(t, continued) + requireOwnedLease(t, ctx, f, false) + record, found, err := f.Store.Execution(). + GetModelCallContext(ctx, testProjectID, f.AgentID, claim.Context.ID) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, executionstore.ModelCallContextFailed, record.State) + require.NotNil(t, record.RetryAt) + require.Equal(t, *record.RetryAt, agentWakeupReadyAt(t, ctx, f.Store, f.AgentID)) + if unfinished { + require.Equal(t, "runtime_released_before_model_result_acceptance", record.ErrorCode) + } + }) + } +} + +func TestAdvanceOwnedAgentWorkFencesInactiveRuntime(t *testing.T) { + t.Parallel() + for _, mode := range []string{"cancel", "expired", "replaced", "wrong project", "wrong runtime"} { + t.Run(mode, func(t *testing.T) { + t.Parallel() + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_fence") + + createToolCallForProcessTest(t, ctx, f, "tool", "read_process") + input := ownedAdvanceInput(f, true) + switch mode { + case "cancel": + _, err := f.Store.Execution().CancelAgent(ctx, executionstore.CancelAgentInput{ + ProjectID: testProjectID, AgentID: f.AgentID, Actor: mustOmnaraActorParams(t, f.UserID), + }) + require.NoError(t, err) + case "expired", "replaced": + expireAgentRuntimeLockForTest(t, ctx, f.Store, f.Lock.ID) + if mode == "replaced" { + n, err := f.Store.Execution().ReapExpiredAgentRuntimeLocks(ctx, 100) + require.NoError(t, err) + require.EqualValues(t, 1, n) + claim, found, err := f.Store.Execution(). + ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + f.Lock = claim.RuntimeLock + } + case "wrong project": + input.ProjectID = uuid.New() + case "wrong runtime": + input.RuntimeLockID = uuid.New() + } + _, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, input) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) + require.False(t, continued) + requireOwnedLease(t, ctx, f, true) + }) + } +} + +func TestAdvanceOwnedAgentWorkReleasesExternalWaitsAndRecoversUnfinishedTools(t *testing.T) { + for _, mode := range []string{"question", "process", "unfinished tool"} { + t.Run(mode, func(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_external_wait") + + toolName := "run_command" + if mode == "question" { + toolName = "ask_question" + } + id := createToolCallForProcessTest(t, ctx, f, "external", toolName) + execute := executionstore.ExecuteToolCallInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, ToolCallID: id, + } + switch mode { + case "question": + createQuestionInteractionForTest(t, ctx, f, id) + case "process": + _, err := startProcessForTest(ctx, f.Store, execute, executionstore.CreateProcessInput{ + ExecutionSpec: processcmd.ForShell("sleep 1", "sh", ""), + AgentMachineBindingID: f.BindingID, Cwd: "/work", + }) + require.NoError(t, err) + case "unfinished tool": + _, err := startAsyncToolCallForTest(ctx, f.Store, execute) + require.NoError(t, err) + } + _, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.False(t, continued) + requireOwnedLease(t, ctx, f, false) + tool, err := f.Store.Execution().GetToolCall(ctx, testProjectID, f.AgentID, id) + require.NoError(t, err) + if mode == "unfinished tool" { + require.Equal(t, executionstore.ToolCallStateCompleted, tool.State) + require.Equal(t, executionstore.ToolResultOutcomeFailed, tool.Outcome) + require.Equal(t, 1, countAgentWakeups(t, ctx, f.Store, f.AgentID)) + } else { + require.Equal(t, executionstore.ToolCallStateWaiting, tool.State) + require.JSONEq(t, `[]`, string(tool.ResultContentParts)) + require.Zero(t, countAgentWakeups(t, ctx, f.Store, f.AgentID)) + } + }) + } +} + +func TestAdvanceOwnedAgentWorkChecksLeaseAfterLockWait(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_expiry_wait") + + createToolCallForProcessTest(t, ctx, f, "tool", "read_process") + blocker, err := f.Store.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = blocker.Rollback(ctx) }() + var pid int32 + require.NoError(t, blocker.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&pid)) + _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, f.AgentID) + require.NoError(t, err) + done := make(chan error, 1) + go func() { + _, _, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + done <- err + }() + integrationdb.WaitForLockWaitBlockedBy(t, ctx, f.Store.pool, "", pid) + _, err = blocker.Exec(ctx, `UPDATE agent_runtime_locks +SET renewed_at=started_at, lease_expires_at=statement_timestamp()-interval '1 microsecond' WHERE id=$1`, + f.Lock.ID) + require.NoError(t, err) + require.NoError(t, blocker.Commit(ctx)) + require.ErrorIs(t, <-done, storeerr.ErrRuntimeLockInactive) + requireOwnedLease(t, ctx, f, true) +} + +func TestAdvanceOwnedAgentWorkCrashRecovery(t *testing.T) { + for _, started := range []bool{false, true} { + name := "before model preparation" + if started { + name = "after model preparation" + } + t.Run(name, func(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_crash") + + id := createToolCallForProcessTest(t, ctx, f, "tool", "read_process") + ownedCompleteTool(t, ctx, f, id) + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + var attempt executionstore.ModelCallClaim + if started { + attempt = ownedModelClaim(t, ctx, f, work) + } + expireAgentRuntimeLockForTest(t, ctx, f.Store, f.Lock.ID) + n, err := f.Store.Execution().ReapExpiredAgentRuntimeLocks(ctx, 100) + require.NoError(t, err) + require.EqualValues(t, 1, n) + replacement, found, err := f.Store.Execution(). + ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + require.NotEqual(t, f.Lock.ID, replacement.RuntimeLock.ID) + require.Equal(t, work.Model.TurnID, replacement.Model.TurnID) + require.Equal(t, work.Model.InputIDs, replacement.Model.InputIDs) + if started { + require.Equal(t, executionstore.ModelWorkResume, replacement.Model.Kind) + require.Equal(t, attempt.Context.ID, replacement.Model.ModelCallContextID) + } else { + require.Equal(t, executionstore.ModelWorkContinue, replacement.Model.Kind) + require.Equal(t, work.Model.SourceModelOutputID, replacement.Model.SourceModelOutputID) + } + _, _, err = f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) + }) + } +} + +func TestAdvanceOwnedAgentWorkCommitFailureRollsBackAdmission(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "owned_rollback") + + input := ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "initial") + var initialTurns int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM agent_turns WHERE agent_id=$1`, f.AgentID, + ).Scan(&initialTurns)) + _, err := f.Store.pool.Exec( + ctx, + `CREATE FUNCTION reject_owned_publication() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN RAISE EXCEPTION 'owned commit rejected'; END $$; +CREATE CONSTRAINT TRIGGER reject_owned_publication AFTER UPDATE ON agent_execution_state +DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION reject_owned_publication();`, + ) + require.NoError(t, err) + _, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.ErrorContains(t, err, "commit advance owned agent work") + require.False(t, continued) + requireOwnedLease(t, ctx, f, true) + var state string + var turns int + require.NoError(t, f.Store.pool.QueryRow(ctx, ` +SELECT state, (SELECT count(*) FROM agent_turns WHERE agent_id=$2) FROM agent_inputs WHERE id=$1`, + input.ID, f.AgentID, + ).Scan(&state, &turns)) + require.Equal(t, "received", state) + require.Equal(t, initialTurns, turns) + require.Zero(t, countAgentWakeups(t, ctx, f.Store, f.AgentID)) + _, err = f.Store.pool.Exec(ctx, `DROP TRIGGER reject_owned_publication ON agent_execution_state`) + require.NoError(t, err) + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + require.Equal(t, []uuid.UUID{input.ID}, work.Model.InputIDs) +} diff --git a/internal/storage/executionstore/agent_work_store.go b/internal/storage/executionstore/agent_work_store.go new file mode 100644 index 000000000..158b761b7 --- /dev/null +++ b/internal/storage/executionstore/agent_work_store.go @@ -0,0 +1,51 @@ +package executionstore + +import ( + "context" + "errors" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type AdvanceOwnedAgentWorkInput struct { + ProjectID uuid.UUID + AgentID uuid.UUID + RuntimeLockID uuid.UUID + AllowModelWork bool + PrepareModel bool +} + +// AdvanceOwnedAgentWork returns continued work or releases the runtime atomically. +// On error, the caller must still release the runtime through the recovery path. +func (s *Store) AdvanceOwnedAgentWork( + ctx context.Context, + input AdvanceOwnedAgentWorkInput, +) (ClaimedAgentWork, bool, error) { + if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil { + return ClaimedAgentWork{}, false, errors.New("project, agent, and runtime lock ids are required") + } + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if errors.Is(err, storeerr.ErrNotFound) { + return ClaimedAgentWork{}, false, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return ClaimedAgentWork{}, false, err + } + defer func() { _ = unit.Rollback(ctx) }() + work, err := h.Advance(ctx, input.RuntimeLockID, input.AllowModelWork, input.PrepareModel) + if err != nil { + return ClaimedAgentWork{}, false, err + } + if err := applyAdmissionDestination(ctx, unit, input.ProjectID, input.AgentID, work.Admission); err != nil { + return ClaimedAgentWork{}, false, err + } + claim, err := shapeOwnedWork(ctx, unit, input.ProjectID, input.AgentID, work) + if err != nil { + return ClaimedAgentWork{}, false, err + } + if err := unit.Commit(ctx, "advance owned agent work"); err != nil { + return ClaimedAgentWork{}, false, err + } + return claim, !work.Released, nil +} diff --git a/internal/storage/executionstore/agents_integration_test.go b/internal/storage/executionstore/agents_integration_test.go index 587bc91f5..8971da126 100644 --- a/internal/storage/executionstore/agents_integration_test.go +++ b/internal/storage/executionstore/agents_integration_test.go @@ -37,17 +37,27 @@ func TestAgentLaunchRequiresConfigAndCanRecordProfile(t *testing.T) { store := newIntegrationStore(pool) user, err := store.Identity().CreateVerifiedUser( ctx, - storagetest.CreateVerifiedUserInput{Email: "agent-profile-launch@example.com", DisplayName: "Agent Profile Launch"}, + storagetest.CreateVerifiedUserInput{ + Email: "agent-profile-launch@example.com", + DisplayName: "Agent Profile Launch", + }, ) if err != nil { t.Fatalf("create user: %v", err) } - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "agent-profile-launch", "Launch Profile", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "agent-profile-launch", + "Launch Profile", + ` instruction: Start with the saved profile config. model: provider_config: openai-prod name: profile-launch -`) +`, + ) launch, err := store.Execution().LaunchAgent(ctx, executionstore.LaunchAgentInput{ ProjectID: testProjectID, @@ -70,9 +80,13 @@ model: t.Fatalf("launched agent state = %s, want active", launch.Agent.State) } requireAgentWakeupCoverage(t, ctx, store, testProjectID, launch.Agent.ID) - if launch.ConfigChange.AgentInput.InputKind != "config_change" || launch.ConfigChange.AgentInput.State != "resolved" || + if launch.ConfigChange.AgentInput.InputKind != "config_change" || + launch.ConfigChange.AgentInput.State != "resolved" || launch.ConfigChange.AgentInput.AgentConfigID != launch.AgentConfig.ID { - t.Fatalf("initial config change should be resolved with launch config: %+v", launch.ConfigChange.AgentInput) + t.Fatalf( + "initial config change should be resolved with launch config: %+v", + launch.ConfigChange.AgentInput, + ) } events, err := store.Execution().ListAgentEventsForRead(ctx, testProjectID, launch.Agent.ID, 0, 10) if err != nil { @@ -117,7 +131,11 @@ model: t.Fatalf("launch config only: %v", err) } if configOnly.Agent.CurrentConfigID != profile.CurrentConfigID { - t.Fatalf("config-only launch current config = %s, want %s", configOnly.Agent.CurrentConfigID, profile.CurrentConfigID) + t.Fatalf( + "config-only launch current config = %s, want %s", + configOnly.Agent.CurrentConfigID, + profile.CurrentConfigID, + ) } if configOnly.Agent.AgentProfileID != uuid.Nil { t.Fatalf("config-only launch profile = %s, want nil", configOnly.Agent.AgentProfileID) @@ -146,12 +164,19 @@ func TestProjectScopedAgentStorageHidesCrossProjectResources(t *testing.T) { ); !errors.Is(err, storeerr.ErrNotFound) { t.Fatalf("list cross-project events error = %v, want ErrNotFound", err) } - if _, _, _, err := store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: otherProjectID, - AgentID: agentID, - ContentBlocks: json.RawMessage(`[{"type":"text","text":"cross-project input"}]`), - IdempotencyKey: "cross-project-input", - }); !errors.Is(err, storeerr.ErrNotFound) { + if _, + _, + _, + err := store.Execution().CreateAgentContentInput(ctx, + executionstore.CreateAgentContentInputInput{ + ProjectID: otherProjectID, + AgentID: agentID, + ContentBlocks: json.RawMessage(`[{"type":"text","text":"cross-project input"}]`), + IdempotencyKey: "cross-project-input", + }); !errors.Is( + err, + storeerr.ErrNotFound, + ) { t.Fatalf("create cross-project agent input error = %v, want ErrNotFound", err) } } @@ -252,13 +277,26 @@ func TestAgentConfigHistoryTablesAreImmutable(t *testing.T) { seedMigratedDB(t, ctx, pool) store := newIntegrationStore(pool) - user := mustCreateProjectDeveloperUser(t, ctx, store, "history-immutable@example.com", "History Immutable") - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "history-immutable", "History Immutable", ` + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "history-immutable@example.com", + "History Immutable", + ) + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "history-immutable", + "History Immutable", + ` instruction: Initial immutable history config. model: provider_config: openai-prod name: history-immutable -`) +`, + ) launch, err := store.Execution().LaunchAgent( ctx, executionstore.LaunchAgentInput{ @@ -406,7 +444,8 @@ model: ); !isPgCode(err, "25006") { t.Fatalf("mutate resolved config-change input rank error = %v, want SQLSTATE 25006", err) } - pinnedConfiguredModel, err := store.Models().GetConfiguredModel(ctx, testOrgID, pinnedConfig.ConfiguredModelID) + pinnedConfiguredModel, err := store.Models(). + GetConfiguredModel(ctx, testOrgID, pinnedConfig.ConfiguredModelID) if err != nil { t.Fatalf("load pinned configured model: %v", err) } @@ -466,7 +505,9 @@ model: if err := store.Execution().DeleteAgentProfile(ctx, testProjectID, profile.ID); err != nil { t.Fatalf("delete profile: %v", err) } - if _, err := store.Execution().GetAgentProfile(ctx, testProjectID, profile.ID); !storeerr.IsNotFound(err) { + if _, err := store.Execution().GetAgentProfile(ctx, testProjectID, profile.ID); !storeerr.IsNotFound( + err, + ) { t.Fatalf("deleted profile lookup error = %v, want not found", err) } var versionCount, deletedVersionCount int @@ -479,7 +520,11 @@ WHERE project_id = $1 AND profile_id = $2 t.Fatalf("count profile versions: %v", err) } if versionCount == 0 || deletedVersionCount != versionCount { - t.Fatalf("profile versions were not soft-deleted: total=%d deleted=%d", versionCount, deletedVersionCount) + t.Fatalf( + "profile versions were not soft-deleted: total=%d deleted=%d", + versionCount, + deletedVersionCount, + ) } if _, err := store.Execution().GetAgentInProject(ctx, testProjectID, launch.Agent.ID); err != nil { t.Fatalf("launched agent should survive profile deletion: %v", err) @@ -489,7 +534,9 @@ WHERE project_id = $1 AND profile_id = $2 ); err != nil || !found { t.Fatalf("profile config should survive profile deletion: %v", err) } - if err := store.Execution().DeleteAgentProfile(ctx, testProjectID, profile.ID); !storeerr.IsNotFound(err) { + if err := store.Execution().DeleteAgentProfile(ctx, testProjectID, profile.ID); !storeerr.IsNotFound( + err, + ) { t.Fatalf("second delete profile error = %v, want not found", err) } } @@ -500,12 +547,19 @@ func TestDeleteAgentProfileSerializesWithRetarget(t *testing.T) { pool := openIntegrationDB(t, ctx) seedMigratedDB(t, ctx, pool) store := newIntegrationStore(pool) - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "delete-retarget", "Delete Retarget", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "delete-retarget", + "Delete Retarget", + ` instruction: Initial profile. model: provider_config: openai-prod name: delete-retarget -`) +`, + ) retargetConfig := mustCreateAgentConfigFromYAML(t, ctx, store, ` instruction: Retargeted profile. model: @@ -759,7 +813,10 @@ model: } duplicateNameInput := input duplicateNameInput.IdempotencyKey = "profile-replay-duplicate-name" - if _, err := store.Execution().CreateAgentProfile(ctx, duplicateNameInput); !errors.Is(err, storeerr.ErrConflict) { + if _, err := store.Execution().CreateAgentProfile(ctx, duplicateNameInput); !errors.Is( + err, + storeerr.ErrConflict, + ) { t.Fatalf("duplicate profile name error = %v, want ErrConflict", err) } retargetYAML := ` @@ -784,17 +841,19 @@ model: if err != nil { t.Fatalf("retarget profile before create replay: %v", err) } - replayedRetarget, err := store.Execution().RetargetAgentProfile(ctx, executionstore.RetargetAgentProfileInput{ - ProjectID: testProjectID, - ProfileID: first.ID, - ExpectedCurrentConfigID: first.CurrentConfigID, - IdempotencyKey: "idem-retarget-profile-replay", - ConfigID: retargetConfig.ID, - }) + replayedRetarget, err := store.Execution(). + RetargetAgentProfile(ctx, executionstore.RetargetAgentProfileInput{ + ProjectID: testProjectID, + ProfileID: first.ID, + ExpectedCurrentConfigID: first.CurrentConfigID, + IdempotencyKey: "idem-retarget-profile-replay", + ConfigID: retargetConfig.ID, + }) if err != nil { t.Fatalf("replay retarget profile: %v", err) } - if replayedRetarget.ID != retargeted.ID || replayedRetarget.CurrentConfigID != retargeted.CurrentConfigID || + if replayedRetarget.ID != retargeted.ID || + replayedRetarget.CurrentConfigID != retargeted.CurrentConfigID || replayedRetarget.CurrentGeneration != retargeted.CurrentGeneration { t.Fatalf("unexpected replayed retarget: retargeted=%+v replayed=%+v", retargeted, replayedRetarget) } @@ -823,7 +882,8 @@ model: if err != nil { t.Fatalf("replay create profile after retarget: %v", err) } - if replayedAfterRetarget.ID != first.ID || replayedAfterRetarget.CurrentConfigID != retargeted.CurrentConfigID || + if replayedAfterRetarget.ID != first.ID || + replayedAfterRetarget.CurrentConfigID != retargeted.CurrentConfigID || replayedAfterRetarget.CurrentConfig.ID != retargeted.CurrentConfigID { t.Fatalf( "replay after retarget should return existing profile head: retargeted=%+v replayed=%+v", @@ -840,7 +900,14 @@ func TestChangeAgentConfigCreatesConfigChangeEventAndIsIdempotent(t *testing.T) seedMigratedDB(t, ctx, pool) store := newIntegrationStore(pool) - user := mustCreateProjectDeveloperUser(t, ctx, store, "agent-config-change@example.com", "Agent Config Change") + + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "agent-config-change@example.com", + "Agent Config Change", + ) otherUser := mustCreateProjectDeveloperUser( t, ctx, @@ -855,12 +922,19 @@ func TestChangeAgentConfigCreatesConfigChangeEventAndIsIdempotent(t *testing.T) "agent-config-change-operator@example.com", "Agent Config Change Operator") - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "agent-config-change", "Change Profile", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "agent-config-change", + "Change Profile", + ` instruction: Original instruction. model: provider_config: openai-prod name: config-change -`) +`, + ) launch, err := store.Execution().LaunchAgent( ctx, executionstore.LaunchAgentInput{ @@ -992,7 +1066,11 @@ model: t.Fatalf("load agent after old replay: %v", err) } if loadedAfterOldReplay.CurrentConfigID != secondChange.AgentConfig.ID { - t.Fatalf("old replay rolled back current config: agent=%+v second=%+v", loadedAfterOldReplay, secondChange) + t.Fatalf( + "old replay rolled back current config: agent=%+v second=%+v", + loadedAfterOldReplay, + secondChange, + ) } conflictInput := changeInputFromRecord(secondChange.AgentConfig) if _, err := store.Execution().ChangeAgentConfig( @@ -1026,12 +1104,19 @@ func TestChangeAgentConfigExpectedCurrentConfigGuard(t *testing.T) { store, "agent-config-expected@example.com", "Agent Config Expected") - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "agent-config-expected", "Expected Profile", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "agent-config-expected", + "Expected Profile", + ` instruction: Original instruction. model: provider_config: openai-prod name: config-expected -`) +`, + ) launch, err := store.Execution().LaunchAgent( ctx, executionstore.LaunchAgentInput{ @@ -1123,17 +1208,27 @@ func TestCaptureAgentConfigForModelContextSeesEventsCommittedBeforeLock(t *testi now := time.Date(2026, 4, 29, 16, 20, 0, 0, time.UTC) user, err := store.Identity().CreateVerifiedUser( ctx, - storagetest.CreateVerifiedUserInput{Email: "agent-config-capture@example.com", DisplayName: "Agent Config Capture"}, + storagetest.CreateVerifiedUserInput{ + Email: "agent-config-capture@example.com", + DisplayName: "Agent Config Capture", + }, ) if err != nil { t.Fatalf("create user: %v", err) } - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "agent-config-capture", "Capture Profile", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "agent-config-capture", + "Capture Profile", + ` instruction: Original instruction. model: provider_config: openai-prod name: config-capture -`) +`, + ) launch, err := store.Execution().LaunchAgent( ctx, executionstore.LaunchAgentInput{ @@ -1173,14 +1268,19 @@ FOR UPDATE resultCh := make(chan captureResult, 1) go func() { close(started) - snapshot, err := store.Execution().CaptureAgentConfigForModelContext(ctx, testProjectID, launch.Agent.ID) + snapshot, err := store.Execution(). + CaptureAgentConfigForModelContext(ctx, testProjectID, launch.Agent.ID) resultCh <- captureResult{snapshot: snapshot, err: err} }() <-started integrationdb.WaitForNamedLockWaiters(t, ctx, pool, "LockAgentInProject", 1) select { case result := <-resultCh: - t.Fatalf("capture returned before the held agent lock was released: snapshot=%+v err=%v", result.snapshot, result.err) + t.Fatalf( + "capture returned before the held agent lock was released: snapshot=%+v err=%v", + result.snapshot, + result.err, + ) default: } @@ -1267,7 +1367,11 @@ WHERE project_id = $3 AND agent_id = $4 AND id = $5 t.Fatalf("capture after lock release: %v", result.err) } if result.snapshot.InputEventSequence != sequence { - t.Fatalf("capture watermark = %d, want committed event sequence %d", result.snapshot.InputEventSequence, sequence) + t.Fatalf( + "capture watermark = %d, want committed event sequence %d", + result.snapshot.InputEventSequence, + sequence, + ) } case <-time.After(5 * time.Second): t.Fatal("capture did not complete after releasing the agent lock") @@ -1282,12 +1386,19 @@ func TestCaptureAgentConfigForEventWatermarkUsesConfigActiveAtSequence(t *testin store := newIntegrationStore(pool) user := mustCreateProjectDeveloperUser(t, ctx, store, "watermark-config@example.com", "Watermark Config") - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "watermark-config", "Watermark Config", ` + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "watermark-config", + "Watermark Config", + ` instruction: First watermark config. model: provider_config: openai-prod name: watermark -`) +`, + ) launch, err := store.Execution().LaunchAgent( ctx, executionstore.LaunchAgentInput{ @@ -1356,13 +1467,26 @@ func TestChangeAgentConfigAcceptsLiveMCPDiffs(t *testing.T) { seedMigratedDB(t, ctx, pool) store := newIntegrationStore(pool) - user := mustCreateProjectDeveloperUser(t, ctx, store, "agent-config-policy@example.com", "Agent Config Policy") - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "agent-config-policy", "Policy Profile", ` + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "agent-config-policy@example.com", + "Agent Config Policy", + ) + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "agent-config-policy", + "Policy Profile", + ` instruction: Original instruction. model: provider_config: openai-prod name: config-policy -`) +`, + ) launch, err := store.Execution().LaunchAgent( ctx, executionstore.LaunchAgentInput{ @@ -1438,7 +1562,13 @@ func TestChangeAgentConfigReconcilesExplicitMachineSources(t *testing.T) { publisher := &recordingPostCommitPublisher{} store := newIntegrationStore(pool, storage.WithPostCommitPublisher(publisher)) now := time.Date(2026, 4, 29, 16, 35, 0, 0, time.UTC) - user := mustCreateProjectDeveloperUser(t, ctx, store, "agent-config-machines@example.com", "Agent Config Machines") + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "agent-config-machines@example.com", + "Agent Config Machines", + ) secret, _, err := store.Secrets().CreateSecret(ctx, secretstore.CreateSecretInput{ OrgID: testOrgID, OwnerKind: secretstore.SecretOwnerProject, @@ -1468,12 +1598,15 @@ func TestChangeAgentConfigReconcilesExplicitMachineSources(t *testing.T) { t.Fatalf("create second machine: %v", err) } for index, machine := range []executionstore.MachineRecord{firstMachine, secondMachine} { - if _, _, err := store.Execution().CreateProjectMachineGrant(ctx, executionstore.CreateProjectMachineGrantInput{ - OrgID: testOrgID, - ProjectID: testProjectID, - MachineID: machine.ID, - IdempotencyKey: fmt.Sprintf("idem-live-explicit-grant-%d", index), - }); err != nil { + if _, + _, + err := store.Execution().CreateProjectMachineGrant(ctx, + executionstore.CreateProjectMachineGrantInput{ + OrgID: testOrgID, + ProjectID: testProjectID, + MachineID: machine.ID, + IdempotencyKey: fmt.Sprintf("idem-live-explicit-grant-%d", index), + }); err != nil { t.Fatalf("grant machine %d: %v", index, err) } } @@ -1521,7 +1654,11 @@ tools: } if shared.MachineBindings[0].MachineID != firstMachine.ID || shared.MachineBindings[0].ID == launch.MachineBindings[0].ID { - t.Fatalf("shared machine bindings = first %+v second %+v", launch.MachineBindings, shared.MachineBindings) + t.Fatalf( + "shared machine bindings = first %+v second %+v", + launch.MachineBindings, + shared.MachineBindings, + ) } invalidYAML := ` instruction: Use explicit machines. @@ -1644,7 +1781,11 @@ tools: } if !sameJSON(firstBinding.EnvOverlay, json.RawMessage(`{"APP":"changed","Base":null,"UNUSED":null}`)) || !sameJSON(firstBinding.SecretEnvOverlay, json.RawMessage(`{"BASE":"`+secret.ID.String()+`"}`)) { - t.Fatalf("reconciled first binding environment = %s / %s", firstBinding.EnvOverlay, firstBinding.SecretEnvOverlay) + t.Fatalf( + "reconciled first binding environment = %s / %s", + firstBinding.EnvOverlay, + firstBinding.SecretEnvOverlay, + ) } reorderedYAML := ` instruction: Use explicit machines. @@ -1676,17 +1817,30 @@ tools: reorderedYAML, "idem-live-explicit-reordered", ) - if reorderedFirst := bindingForMachine(firstMachine.ID); !reorderedFirst.UpdatedAt.Equal(firstBinding.UpdatedAt) { - t.Fatalf("source reorder updated first binding at %v, want %v", reorderedFirst.UpdatedAt, firstBinding.UpdatedAt) + if reorderedFirst := bindingForMachine(firstMachine.ID); !reorderedFirst.UpdatedAt.Equal( + firstBinding.UpdatedAt, + ) { + t.Fatalf( + "source reorder updated first binding at %v, want %v", + reorderedFirst.UpdatedAt, + firstBinding.UpdatedAt, + ) } - if reorderedSecond := bindingForMachine(secondMachine.ID); !reorderedSecond.UpdatedAt.Equal(secondBinding.UpdatedAt) { - t.Fatalf("source reorder updated second binding at %v, want %v", reorderedSecond.UpdatedAt, secondBinding.UpdatedAt) + if reorderedSecond := bindingForMachine(secondMachine.ID); !reorderedSecond.UpdatedAt.Equal( + secondBinding.UpdatedAt, + ) { + t.Fatalf( + "source reorder updated second binding at %v, want %v", + reorderedSecond.UpdatedAt, + secondBinding.UpdatedAt, + ) } - token, err := store.Execution().CreateBYOMachineDaemonToken(ctx, executionstore.CreateBYOMachineDaemonTokenInput{ - OrgID: testOrgID, - MachineID: firstMachine.ID, - Name: "live source removal", - }) + token, err := store.Execution(). + CreateBYOMachineDaemonToken(ctx, executionstore.CreateBYOMachineDaemonTokenInput{ + OrgID: testOrgID, + MachineID: firstMachine.ID, + Name: "live source removal", + }) if err != nil { t.Fatalf("create live source removal daemon token: %v", err) } @@ -1766,7 +1920,12 @@ tools: markProcessStartedForTest(t, ctx, fixture, process, processAt.Add(200*time.Millisecond)) return process } - firstProcess := startRunningProcess("live_explicit_source_removed", launch.Agent.ID, firstBinding.ID, firstLock) + firstProcess := startRunningProcess( + "live_explicit_source_removed", + launch.Agent.ID, + firstBinding.ID, + firstLock, + ) sharedProcess := startRunningProcess( "live_explicit_source_shared", shared.Agent.ID, @@ -1788,7 +1947,6 @@ tools: t, ctx, store, - secondOnlyYAML, ) // A source removal must lock only its own agent, even when another agent @@ -1802,13 +1960,15 @@ tools: } removalCtx, cancelRemoval := context.WithTimeout(ctx, 2*time.Second) defer cancelRemoval() - if _, err := store.Execution().IntegrationChangeAgentConfigOnce(removalCtx, executionstore.ChangeAgentConfigInput{ - CreateAgentConfigInput: changeInputFromRecord(removalConfig), - AgentID: launch.Agent.ID, - ActorType: identitystore.PrincipalTypeUser, - ActorID: user.ID, - IdempotencyKey: "idem-live-explicit-remove", - }); err != nil { + if _, + err := store.Execution().IntegrationChangeAgentConfigOnce(removalCtx, + executionstore.ChangeAgentConfigInput{ + CreateAgentConfigInput: changeInputFromRecord(removalConfig), + AgentID: launch.Agent.ID, + ActorType: identitystore.PrincipalTypeUser, + ActorID: user.ID, + IdempotencyKey: "idem-live-explicit-remove", + }); err != nil { t.Fatalf("remove explicit source while another grant agent is locked: %v", err) } if err := sharedTx.Commit(ctx); err != nil { @@ -1831,7 +1991,8 @@ tools: if !publisher.hasProcessTermination(firstMachine.ID, firstProcess.ID) { t.Fatalf("removed explicit source did not publish termination for process %s", firstProcess.ID) } - remainingProcess, err := store.Execution().GetProcess(ctx, testProjectID, shared.Agent.ID, sharedProcess.ID) + remainingProcess, err := store.Execution(). + GetProcess(ctx, testProjectID, shared.Agent.ID, sharedProcess.ID) if err != nil { t.Fatalf("load shared explicit source process: %v", err) } @@ -1904,7 +2065,8 @@ instruction: test ` compiledA := mustCompileAgentYAMLResolved(t, ctx, store, sourceA) compiledB := mustCompileAgentYAMLResolved(t, ctx, store, sourceB) - if compiledA.Hash != compiledB.Hash || string(compiledA.CanonicalJSON) != string(compiledB.CanonicalJSON) { + if compiledA.Hash != compiledB.Hash || + string(compiledA.CanonicalJSON) != string(compiledB.CanonicalJSON) { t.Fatalf( "test sources should compile to equivalent config: hash %q/%q json %s/%s", compiledA.Hash, @@ -1935,7 +2097,10 @@ instruction: test t.Fatalf("create equivalent source B config: %v", err) } if equivalentA.ID == equivalentB.ID { - t.Fatalf("same behavior with different authored source should keep distinct config rows: %s", equivalentA.ID) + t.Fatalf( + "same behavior with different authored source should keep distinct config rows: %s", + equivalentA.ID, + ) } if equivalentA.EffectiveDefinitionHash != equivalentB.EffectiveDefinitionHash { t.Fatalf( @@ -1999,7 +2164,11 @@ func mustCreateAgentConfigFromYAML( func mustCompileAgentYAML(t *testing.T, sourceYAML string) agentconfig.Result { t.Helper() - compiled, err := agentconfig.Compile(agentconfig.SourceFormatYAML, []byte(sourceYAML), agentconfig.CompileOptions{}) + compiled, err := agentconfig.Compile( + agentconfig.SourceFormatYAML, + []byte(sourceYAML), + agentconfig.CompileOptions{}, + ) if err != nil { t.Fatalf("compile agent yaml: %v", err) } @@ -2056,7 +2225,14 @@ func TestEventWebhookTargetFollowsCurrentConfig(t *testing.T) { store := newIntegrationStore(pool) user := mustCreateProjectDeveloperUser(t, ctx, store, "event-webhook@example.com", "Event webhook") source := "instruction: Test event webhook.\nmodel:\n provider_config: openai-prod\n name: event-webhook\n" - profile := mustCreateConfigAndProfileBookmarkFromYAML(t, ctx, store, "event-webhook", "Event webhook", source) + profile := mustCreateConfigAndProfileBookmarkFromYAML( + t, + ctx, + store, + "event-webhook", + "Event webhook", + source, + ) launch, err := store.Execution().LaunchAgent(ctx, executionstore.LaunchAgentInput{ ProjectID: testProjectID, ProfileID: profile.ID, AgentConfigID: profile.CurrentConfigID, LaunchedBy: userPrincipal(user.ID), IdempotencyKey: "event-webhook-launch", @@ -2074,10 +2250,16 @@ func TestEventWebhookTargetFollowsCurrentConfig(t *testing.T) { _, err := store.Execution().ChangeAgentConfig(ctx, executionstore.ChangeAgentConfigInput{ CreateAgentConfigInput: executionstore.CreateAgentConfigInput{ ProjectID: testProjectID, Source: updated, SourceFormat: "yaml", - ConfiguredModelID: parseConfiguredModelID(t, compiled), CompiledDefinition: compiled.CanonicalJSON, + ConfiguredModelID: parseConfiguredModelID( + t, + compiled, + ), CompiledDefinition: compiled.CanonicalJSON, EffectiveDefinitionHash: compiled.Hash, }, - AgentID: launch.Agent.ID, ActorType: identitystore.PrincipalTypeUser, ActorID: user.ID, Reason: "user_update", + AgentID: launch.Agent.ID, + ActorType: identitystore.PrincipalTypeUser, + ActorID: user.ID, + Reason: "user_update", }) require.NoError(t, err) target, err := store.Execution().GetAgentEventWebhookTarget(ctx, launch.Agent.ID) diff --git a/internal/storage/executionstore/context_checkpoints_boundary.go b/internal/storage/executionstore/context_checkpoints_boundary.go deleted file mode 100644 index 7ab1f3d77..000000000 --- a/internal/storage/executionstore/context_checkpoints_boundary.go +++ /dev/null @@ -1,71 +0,0 @@ -package executionstore - -import ( - "context" - "errors" - "fmt" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" -) - -var ErrCheckpointBoundaryUnsafe = errors.New("context checkpoint boundary is unsafe") - -func validateClosedCheckpointRangeTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID uuid.UUID, - start, end int64, -) error { - count, err := dbsqlc.New(tx).CountCheckpointRangeEvents( - ctx, - dbsqlc.CountCheckpointRangeEventsParams{ - ProjectID: projectID, - AgentID: agentID, - StartSequence: start, - EndSequence: end, - }, - ) - if err != nil { - return fmt.Errorf("validate checkpoint range: %w", err) - } - if count != end-start+1 { - return fmt.Errorf( - "%w: range %d..%d is not a closed event range", - ErrCheckpointBoundaryUnsafe, - start, - end, - ) - } - return nil -} - -func validateCheckpointDoesNotCutOpenAuthoritiesTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID uuid.UUID, - start, end int64, -) error { - count, err := dbsqlc.New(tx).CountCheckpointRangeOpenToolResults( - ctx, - dbsqlc.CountCheckpointRangeOpenToolResultsParams{ - ProjectID: projectID, - AgentID: agentID, - StartSequence: start, - EndSequence: end, - }, - ) - if err != nil { - return fmt.Errorf("validate checkpoint open tool results: %w", err) - } - if count != 0 { - return fmt.Errorf( - "%w: range %d..%d cuts an active tool call", - ErrCheckpointBoundaryUnsafe, - start, - end, - ) - } - return nil -} diff --git a/internal/storage/executionstore/context_checkpoints_store.go b/internal/storage/executionstore/context_checkpoints_store.go index 3e8039f2b..80851eccb 100644 --- a/internal/storage/executionstore/context_checkpoints_store.go +++ b/internal/storage/executionstore/context_checkpoints_store.go @@ -2,15 +2,11 @@ package executionstore import ( "context" - "encoding/json" - "errors" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/modelprotocol" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -34,203 +30,37 @@ func (s *Store) PublishContextCheckpoint( ctx context.Context, input PublishContextCheckpointInput, ) (ContextCheckpointRecord, error) { - if err := validatePublishContextCheckpointInput(input); err != nil { - return ContextCheckpointRecord{}, err - } - input.Usage = modelUsageForStorage(input.Usage) - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("begin publish context checkpoint: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if _, err := q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, - ID: input.AgentID, - }); err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("lock agent for context checkpoint: %w", err) - } - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return ContextCheckpointRecord{}, err } - contextRow, err := loadModelCallContextByID( + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.PublishCheckpoint( ctx, - q, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, + agentexecution.PublishCheckpointInput{RuntimeLockID: input.RuntimeLockID, + ContextID: input.ModelCallContextID, Summary: input.Summary, Evidence: agentexecution.ModelEvidence{ + APIFormat: input.APIFormat, APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, + ResponseID: input.ProviderResponseID, + Usage: input.Usage, + Cost: input.ProviderReportedCostUSD, + Metadata: input.ProviderMetadata}}, ) - if err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("load compaction context: %w", err) - } - if contextRow.OperationKind != ModelCallOperationCompaction || - contextRow.State != ModelCallContextStarted || - contextRow.RuntimeLockID != input.RuntimeLockID || - contextRow.SourceEventSequenceEnd == nil { - return ContextCheckpointRecord{}, storeerr.ErrStateTransitionConflict - } - sourceStart, err := compactionSourceStartTx(ctx, q, contextRow) if err != nil { return ContextCheckpointRecord{}, err } - if err := validateClosedCheckpointRangeTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - sourceStart, - *contextRow.SourceEventSequenceEnd, - ); err != nil { - return ContextCheckpointRecord{}, err - } - if err := validateCheckpointDoesNotCutOpenAuthoritiesTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - sourceStart, - *contextRow.SourceEventSequenceEnd, - ); err != nil { - return ContextCheckpointRecord{}, err - } - - checkpointID, err := uuid.NewV7() - if err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("generate checkpoint id: %w", err) - } - checkpointEventID, err := uuid.NewV7() - if err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("generate checkpoint event id: %w", err) - } - turnID, err := q.GetModelCallContextTurnID(ctx, dbsqlc.GetModelCallContextTurnIDParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - }) - if err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("load checkpoint turn: %w", err) - } - _, err = q.InsertContextCheckpoint(ctx, dbsqlc.InsertContextCheckpointParams{ - ID: checkpointID, - Summary: input.Summary, - ProducerModelCallContextID: input.ModelCallContextID, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }) - if errors.Is(err, pgx.ErrNoRows) { - return ContextCheckpointRecord{}, storeerr.ErrRuntimeLockInactive - } - if err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("insert context checkpoint: %w", err) - } - eventRecord, err := appendTypedAgentEventTx( - ctx, - txNotifications, - tx, - AppendTypedAgentEventInput{ - ID: checkpointEventID, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: turnID, - Kind: events.KindContextCheckpoint, - ContextCheckpointID: checkpointID, - }, - ) - if err != nil { - return ContextCheckpointRecord{}, err - } - if err := updateAgentTurnLatestEventTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - turnID, - eventRecord.Event.ID, - uuid.Nil, - ); err != nil { - return ContextCheckpointRecord{}, err - } - if _, err := finishModelCallContextTx(ctx, q, finishModelCallContextInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ToState: ModelCallContextSucceeded, - APIFormat: input.APIFormat, - APIVariant: input.APIVariant, - ProviderRequestID: input.ProviderRequestID, - ProviderResponseID: input.ProviderResponseID, - Usage: input.Usage, - ProviderReportedCostUSD: input.ProviderReportedCostUSD, - ProviderMetadata: input.ProviderMetadata, - }); err != nil { - return ContextCheckpointRecord{}, err - } - if err := q.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - Metadata: json.RawMessage(`{"reason":"context_checkpoint"}`), - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }); err != nil { - return ContextCheckpointRecord{}, fmt.Errorf("reconcile wakeup after context checkpoint: %w", err) + if !result.Created { + return ContextCheckpointRecord{}, storeerr.ErrStateTransitionConflict } - record, found, err := getContextCheckpointByProducerContextTx( - ctx, - q, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) + row, err := dbsqlc.New(unit.DB()). + GetContextCheckpoint(ctx, + dbsqlc.GetContextCheckpointParams{ProjectID: input.ProjectID, + AgentID: input.AgentID, + ID: result.ID}) if err != nil { return ContextCheckpointRecord{}, err } - if !found { - return ContextCheckpointRecord{}, storeerr.ErrStateTransitionConflict - } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "publish context checkpoint"); err != nil { + if err = unit.Commit(ctx, "publish context checkpoint"); err != nil { return ContextCheckpointRecord{}, err } - return record, nil -} - -func compactionSourceStartTx( - ctx context.Context, - q *dbsqlc.Queries, - contextRow ModelCallContextRecord, -) (int64, error) { - if contextRow.OperationKind != ModelCallOperationCompaction { - return 0, storeerr.ErrStateTransitionConflict - } - checkpoint, err := q.GetLatestApplicableContextCheckpoint( - ctx, - dbsqlc.GetLatestApplicableContextCheckpointParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - MaxEventSequence: contextRow.InputEventSequence, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return 1, nil - } - if err != nil { - return 0, fmt.Errorf("load prior compaction checkpoint: %w", err) - } - return checkpoint.SummarizedThroughEventSequence + 1, nil -} - -func validatePublishContextCheckpointInput(input PublishContextCheckpointInput) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.ModelCallContextID == uuid.Nil || input.Summary == "" { - return errors.New("project, agent, runtime, context, and summary are required") - } - if err := modelenvelope.ValidateProviderReportedCostUSD(input.ProviderReportedCostUSD); err != nil { - return fmt.Errorf("provider-reported cost: %w", err) - } - return nil + return contextCheckpointRecordFromGetSQLC(row), nil } diff --git a/internal/storage/executionstore/context_checkpoints_terminal_failure.go b/internal/storage/executionstore/context_checkpoints_terminal_failure.go index ae9c82898..b8c7d9c28 100644 --- a/internal/storage/executionstore/context_checkpoints_terminal_failure.go +++ b/internal/storage/executionstore/context_checkpoints_terminal_failure.go @@ -3,15 +3,11 @@ package executionstore import ( "context" "encoding/json" - "errors" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/modelprotocol" - "github.com/omnara-ai/omnara/internal/notifications" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) type RecordTerminalCompactionFailureInput struct { @@ -37,78 +33,25 @@ func (s *Store) RecordTerminalCompactionFailure( ctx context.Context, input RecordTerminalCompactionFailureInput, ) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.ModelCallContextID == uuid.Nil || - input.ErrorKind == "" || input.ErrorMessage == "" { - return errors.New("project, agent, runtime, compaction context, and error are required") - } - var err error - input.ErrorDetails, err = normalizedJSONObject(input.ErrorDetails, "compaction error details") - if err != nil { - return err - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return fmt.Errorf("begin terminal compaction failure: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return err } - if err := recordTerminalCompactionFailureTx( + defer func() { _ = unit.Rollback(ctx) }() + _, err = h.FailModel( ctx, - txNotifications, - tx, - q, - input, - modelCallContextRuntimeOwned, - ); err != nil { + agentexecution.ModelFailure{ContextID: input.ModelCallContextID, RuntimeLockID: input.RuntimeLockID, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + ServedModel: input.ServedProviderModelSlug, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, + APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, + Usage: input.Usage, Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}}, + ) + if err != nil { return err } - return s.commitTxWithNotifications(ctx, tx, txNotifications, "record terminal compaction failure") -} - -func recordTerminalCompactionFailureTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - q *dbsqlc.Queries, - input RecordTerminalCompactionFailureInput, - runtimeAuthority modelCallContextRuntimeAuthority, -) error { - modelFailure := RecordModelCallErrorAndCompleteContextInput{} - modelFailure.ProjectID = input.ProjectID - modelFailure.AgentID = input.AgentID - modelFailure.RuntimeLockID = input.RuntimeLockID - modelFailure.ModelCallContextID = input.ModelCallContextID - modelFailure.APIFormat = input.APIFormat - modelFailure.APIVariant = input.APIVariant - modelFailure.ServedProviderModelSlug = input.ServedProviderModelSlug - modelFailure.ProviderRequestID = input.ProviderRequestID - modelFailure.ProviderResponseID = input.ProviderResponseID - modelFailure.ErrorKind = input.ErrorKind - modelFailure.ErrorCode = input.ErrorCode - modelFailure.ErrorMessage = input.ErrorMessage - modelFailure.ErrorDetails = input.ErrorDetails - modelFailure.Usage = input.Usage - modelFailure.ProviderReportedCostUSD = input.ProviderReportedCostUSD - modelFailure.ProviderMetadata = input.ProviderMetadata - _, err := recordTerminalModelCallFailureTx( - ctx, - txNotifications, - tx, - q, - modelFailure, - runtimeAuthority, - ModelCallOperationCompaction, - ) - return err + return unit.Commit(ctx, "record terminal compaction failure") } diff --git a/internal/storage/executionstore/context_events_store.go b/internal/storage/executionstore/context_events_store.go index 1bf1a42d4..cb2380e6b 100644 --- a/internal/storage/executionstore/context_events_store.go +++ b/internal/storage/executionstore/context_events_store.go @@ -65,6 +65,10 @@ func (s *Store) ListContextEvents( if err != nil { return nil, fmt.Errorf("list context events: %w", err) } + return contextEventsFromSQLC(rows, projectID, agentID), nil +} + +func contextEventsFromSQLC(rows []dbsqlc.ListContextEventsRow, projectID, agentID uuid.UUID) []ContextEventRecord { out := make([]ContextEventRecord, 0, len(rows)) for _, row := range rows { record := ContextEventRecord{ @@ -92,7 +96,7 @@ func (s *Store) ListContextEvents( } out = append(out, record) } - return out, nil + return out } func (s *Store) IsOutputLimitBoundary(ctx context.Context, projectID, agentID uuid.UUID, sequence int64) (bool, error) { diff --git a/internal/storage/executionstore/cron_trigger_agent_inputs_store.go b/internal/storage/executionstore/cron_trigger_agent_inputs_store.go index 5707d2f45..9bed4ffe9 100644 --- a/internal/storage/executionstore/cron_trigger_agent_inputs_store.go +++ b/internal/storage/executionstore/cron_trigger_agent_inputs_store.go @@ -7,7 +7,7 @@ import ( "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -51,16 +51,17 @@ func (s *Store) CreateCronTriggerAgentInput(ctx context.Context, input CreateCro return err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return fmt.Errorf("begin cron trigger input delivery: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + qtx := dbsqlc.New(tx) // Input creation locks the agent before completion locks the trigger, matching // the lock order used when archiving an agent and deleting its cron triggers. - if _, err := createAgentContentInputTx(ctx, txNotifications, tx, qtx, contentInput, contentBlocks); err != nil { + if _, err := createAgentContentInputTx(ctx, unit, qtx, contentInput, contentBlocks); err != nil { return err } if err := completeCronTriggerFiringTx(ctx, qtx, CompleteCronTriggerFiringInput{ @@ -71,5 +72,5 @@ func (s *Store) CreateCronTriggerAgentInput(ctx context.Context, input CreateCro }); err != nil { return err } - return s.commitTxWithNotifications(ctx, tx, txNotifications, "deliver cron trigger input") + return unit.Commit(ctx, "deliver cron trigger input") } diff --git a/internal/storage/executionstore/db_values.go b/internal/storage/executionstore/db_values.go index e94cb1c02..0f1b5f9f6 100644 --- a/internal/storage/executionstore/db_values.go +++ b/internal/storage/executionstore/db_values.go @@ -20,13 +20,6 @@ func rawMessageFromSQLCPtr(value *json.RawMessage) json.RawMessage { return *value } -func sqlcRawMessageFromEmpty(value json.RawMessage) *json.RawMessage { - if len(value) == 0 { - return nil - } - return &value -} - func stringFromSQLCText(value *string) string { if value == nil { return "" diff --git a/internal/storage/executionstore/default_reconciliation.go b/internal/storage/executionstore/default_reconciliation.go index fe4691ea8..befb22524 100644 --- a/internal/storage/executionstore/default_reconciliation.go +++ b/internal/storage/executionstore/default_reconciliation.go @@ -22,13 +22,13 @@ type defaultMachinePoolReconciliation struct { func (s *Store) ReconcileDefaultMachinePoolsTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, orgID uuid.UUID, templates []DefaultMachinePoolTemplate, rows []dbsqlc.MachinePool, apply bool, ) ([]string, error) { - qtx := s.q.WithTx(tx) + qtx := dbsqlc.New(tx) targets := make([]defaultMachinePoolReconciliation, 0, len(rows)) for _, template := range templates { name, err := resourcename.CanonicalizeRequired("machine pool name", template.Name) diff --git a/internal/storage/executionstore/event_webhooks_integration_test.go b/internal/storage/executionstore/event_webhooks_integration_test.go index 07add13c5..e75382c11 100644 --- a/internal/storage/executionstore/event_webhooks_integration_test.go +++ b/internal/storage/executionstore/event_webhooks_integration_test.go @@ -28,29 +28,55 @@ func enableEventWebhook(t *testing.T, fixture processDaemonFixture, events []str source := "instruction: Webhook test.\nmodel:\n provider_config: openai-prod\n name: webhook-test\n" + "event_webhook:\n url: https://example.com/events\n" if events == nil { - events = []string{"agent_input", "model_output", "tool_result", "context_checkpoint", "tool_call_update"} + events = []string{ + "agent_input", + "model_output", + "tool_result", + "context_checkpoint", + "tool_call_update", + } } encoded, err := json.Marshal(events) require.NoError(t, err) source += " events: " + string(encoded) + "\n" - user := mustCreateProjectDeveloperUser(t, t.Context(), fixture.Store, "webhook-config@example.com", "Webhook") + user := mustCreateProjectDeveloperUser( + t, + t.Context(), + fixture.Store, + "webhook-config@example.com", + "Webhook", + ) compiled := mustCompileAgentYAMLResolved(t, t.Context(), fixture.Store, source) _, err = fixture.Store.Execution().ChangeAgentConfig(t.Context(), executionstore.ChangeAgentConfigInput{ CreateAgentConfigInput: executionstore.CreateAgentConfigInput{ ProjectID: testProjectID, Source: source, SourceFormat: "yaml", - ConfiguredModelID: parseConfiguredModelID(t, compiled), CompiledDefinition: compiled.CanonicalJSON, + ConfiguredModelID: parseConfiguredModelID( + t, + compiled, + ), CompiledDefinition: compiled.CanonicalJSON, EffectiveDefinitionHash: compiled.Hash, }, - AgentID: fixture.AgentID, ActorType: identitystore.PrincipalTypeUser, ActorID: user.ID, Reason: "user_update", + AgentID: fixture.AgentID, + ActorType: identitystore.PrincipalTypeUser, + ActorID: user.ID, + Reason: "user_update", }) require.NoError(t, err) } -func assertWebhookToolStates(t *testing.T, fixture processDaemonFixture, toolID uuid.UUID, expected []string) { +func assertWebhookToolStates( + t *testing.T, + fixture processDaemonFixture, + toolID uuid.UUID, + expected []string, +) { t.Helper() - rows, err := fixture.Store.pool.Query(t.Context(), + rows, err := fixture.Store.pool.Query( + t.Context(), "SELECT tool_state FROM event_webhook_deliveries WHERE agent_id = $1 AND tool_call_id = $2 ORDER BY id", - fixture.AgentID, toolID) + fixture.AgentID, + toolID, + ) require.NoError(t, err) defer rows.Close() var states []string @@ -109,7 +135,10 @@ func TestInteractionWebhookLifecycle(t *testing.T) { "interaction_webhook", "ask_question", kind == "question") enableEventWebhook(t, fixture, []string{"tool_call_update"}) publisher := &recordingPostCommitPublisher{} - fixture.Store = newIntegrationStore(fixture.Store.pool, storage.WithPostCommitPublisher(publisher)) + fixture.Store = newIntegrationStore( + fixture.Store.pool, + storage.WithPostCommitPublisher(publisher), + ) var interaction executionstore.AgentInteractionRecord openingState := "awaiting_permission" if kind == "permission" { @@ -132,8 +161,10 @@ func TestInteractionWebhookLifecycle(t *testing.T) { } input := executionstore.ResolveAgentInteractionInput{ ProjectID: testProjectID, AgentID: fixture.AgentID, ID: interaction.ID, - Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{OptionIndices: []int{option}}}}, - Actor: mustOmnaraActorParams(t, fixture.UserID), + Resolution: interactionform.Resolution{ + Answers: []interactionform.Answer{{OptionIndices: []int{option}}}, + }, + Actor: mustOmnaraActorParams(t, fixture.UserID), } for range 2 { _, err := fixture.Store.Execution().ResolveAgentInteraction(ctx, input) @@ -151,30 +182,38 @@ func TestInteractionWebhookLifecycle(t *testing.T) { ) require.NoError(t, err) case "steer": - _, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"change direction"}]`), - DeliveryMode: executionstore.DeliveryModeSteering, CancelOpenInteractions: true, - IdempotencyKey: "cancel-interaction", - }) + _, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"change direction"}]`), + DeliveryMode: executionstore.DeliveryModeSteering, CancelOpenInteractions: true, + IdempotencyKey: "cancel-interaction", + }) require.NoError(t, err) case "fail": - _, err := fixture.Store.Execution().CompleteToolCall(ctx, executionstore.CompleteToolCallInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, ID: toolID, - RuntimeLockID: fixture.Lock.ID, Outcome: executionstore.ToolResultOutcomeFailed, - }) + _, err := fixture.Store.Execution(). + CompleteToolCall(ctx, executionstore.CompleteToolCallInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, ID: toolID, + RuntimeLockID: fixture.Lock.ID, Outcome: executionstore.ToolResultOutcomeFailed, + }) require.NoError(t, err) } - assertInteractionWebhookUpdates(t, fixture, publisher, toolID, []notifications.ToolCallUpdatedCommitted{ - {AgentID: fixture.AgentID, ToolCallID: toolID, State: openingState, - InteractionUpdate: ¬ifications.AgentInteractionUpdate{ - ID: interaction.ID, InteractionKind: kind, State: "open", - }}, - {AgentID: fixture.AgentID, ToolCallID: toolID, State: closingState, - InteractionUpdate: ¬ifications.AgentInteractionUpdate{ - ID: interaction.ID, InteractionKind: kind, State: interactionState, - }}, - }) + assertInteractionWebhookUpdates( + t, + fixture, + publisher, + toolID, + []notifications.ToolCallUpdatedCommitted{ + {AgentID: fixture.AgentID, ToolCallID: toolID, State: openingState, + InteractionUpdate: ¬ifications.AgentInteractionUpdate{ + ID: interaction.ID, InteractionKind: kind, State: "open", + }}, + {AgentID: fixture.AgentID, ToolCallID: toolID, State: closingState, + InteractionUpdate: ¬ifications.AgentInteractionUpdate{ + ID: interaction.ID, InteractionKind: kind, State: interactionState, + }}, + }, + ) }) } } @@ -213,7 +252,13 @@ func TestInteractionWebhookCancelMultipleTools(t *testing.T) { ID: interaction.ID, InteractionKind: string(interaction.InteractionKind), State: "canceled", } } - assertInteractionWebhookUpdates(t, fixture, publisher, toolIDs[i], []notifications.ToolCallUpdatedCommitted{expected}) + assertInteractionWebhookUpdates( + t, + fixture, + publisher, + toolIDs[i], + []notifications.ToolCallUpdatedCommitted{expected}, + ) } } @@ -221,6 +266,7 @@ func TestInteractionWebhookPermissionThenQuestion(t *testing.T) { t.Parallel() ctx := t.Context() fixture := newProcessDaemonFixture(t, ctx, "permission_then_question_webhook") + toolID := createToolCallForProcessTestWithPermission(t, ctx, fixture, "permission_then_question_webhook", "ask_question", false) enableEventWebhook(t, fixture, []string{"tool_call_update"}) @@ -276,17 +322,22 @@ func TestEventWebhookFiltersBeforeEnqueue(t *testing.T) { ctx := t.Context() fixture := newProcessDaemonFixture(t, ctx, "webhook_filter") enableEventWebhook(t, fixture, scenario.events) - _, err := fixture.Store.pool.Exec(ctx, "DELETE FROM event_webhook_deliveries WHERE agent_id = $1", fixture.AgentID) + _, err := fixture.Store.pool.Exec( + ctx, + "DELETE FROM event_webhook_deliveries WHERE agent_id = $1", + fixture.AgentID, + ) require.NoError(t, err) pending := notifications.NewTxNotifications() for i, kind := range []string{"agent_input", "model_output", "tool_result", "context_checkpoint"} { pending.AddAgentEvent(fixture.AgentID, int64(i+1), kind) } pending.AddToolCallUpdate(fixture.AgentID, uuid.New(), "ready", nil) - tx, err := fixture.Store.pool.Begin(ctx) + unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) - defer func() { _ = tx.Rollback(ctx) }() - require.NoError(t, fixture.Store.Execution().IntegrationCommitTxWithNotifications(ctx, tx, pending, "filter test")) + defer func() { _ = unit.Rollback(ctx) }() + unit.Notifications().Restore(pending) + require.NoError(t, unit.Commit(ctx, "filter test")) var sequences []int64 var toolUpdates int require.NoError(t, fixture.Store.pool.QueryRow(ctx, ` @@ -307,16 +358,25 @@ func TestEventWebhookClaimsRecoverAndExpire(t *testing.T) { ctx := t.Context() fixture := newProcessDaemonFixture(t, ctx, "webhook_claims") enableEventWebhook(t, fixture, nil) - _, err := fixture.Store.pool.Exec(ctx, "DELETE FROM event_webhook_deliveries WHERE agent_id = $1", fixture.AgentID) + _, err := fixture.Store.pool.Exec( + ctx, + "DELETE FROM event_webhook_deliveries WHERE agent_id = $1", + fixture.AgentID, + ) require.NoError(t, err) pending := notifications.NewTxNotifications() toolID := uuid.New() - interaction := notifications.AgentInteractionUpdate{ID: uuid.New(), InteractionKind: "permission", State: "resolved"} + interaction := notifications.AgentInteractionUpdate{ + ID: uuid.New(), + InteractionKind: "permission", + State: "resolved", + } pending.AddToolCallUpdate(fixture.AgentID, toolID, "ready", &interaction) - tx, err := fixture.Store.pool.Begin(ctx) + unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) - defer func() { _ = tx.Rollback(ctx) }() - require.NoError(t, fixture.Store.Execution().IntegrationCommitTxWithNotifications(ctx, tx, pending, "webhook test")) + defer func() { _ = unit.Rollback(ctx) }() + unit.Notifications().Restore(pending) + require.NoError(t, unit.Commit(ctx, "webhook test")) var attempts int var claimed bool require.NoError(t, fixture.Store.pool.QueryRow(ctx, ` @@ -332,18 +392,25 @@ func TestEventWebhookClaimsRecoverAndExpire(t *testing.T) { require.Equal(t, int32(1), first.AttemptCount) _, err = fixture.Store.Execution().ClaimEventWebhookDelivery(ctx, 128) require.ErrorIs(t, err, storeerr.ErrNotFound) - _, err = fixture.Store.pool.Exec(ctx, - "UPDATE event_webhook_deliveries SET claim_expires_at = statement_timestamp() - interval '1 second' WHERE id = $1", first.ID) + _, err = fixture.Store.pool.Exec( + ctx, + "UPDATE event_webhook_deliveries SET claim_expires_at = statement_timestamp() - interval '1 second' WHERE id = $1", + first.ID, + ) require.NoError(t, err) second, err := fixture.Store.Execution().ClaimEventWebhookDelivery(ctx, 128) require.NoError(t, err) require.Equal(t, first.ID, second.ID) require.NotEqual(t, first.ClaimToken, second.ClaimToken) require.Equal(t, int32(2), second.AttemptCount) - require.NoError(t, fixture.Store.Execution().CompleteEventWebhookDelivery(ctx, first.ID, first.ClaimToken)) + require.NoError( + t, + fixture.Store.Execution().CompleteEventWebhookDelivery(ctx, first.ID, first.ClaimToken), + ) _, err = fixture.Store.Execution().RetryEventWebhookDelivery(ctx, first.ID, first.ClaimToken, time.Minute) require.ErrorIs(t, err, storeerr.ErrNotFound) - retry, err := fixture.Store.Execution().RetryEventWebhookDelivery(ctx, second.ID, second.ClaimToken, time.Minute) + retry, err := fixture.Store.Execution(). + RetryEventWebhookDelivery(ctx, second.ID, second.ClaimToken, time.Minute) require.NoError(t, err) require.False(t, retry.GaveUp) var nextAttemptAt time.Time @@ -352,19 +419,25 @@ func TestEventWebhookClaimsRecoverAndExpire(t *testing.T) { require.True(t, nextAttemptAt.Equal(retry.NextAttemptAt)) _, err = fixture.Store.Execution().ClaimEventWebhookDelivery(ctx, 128) require.ErrorIs(t, err, storeerr.ErrNotFound) - _, err = fixture.Store.pool.Exec(ctx, - "UPDATE event_webhook_deliveries SET next_attempt_at = statement_timestamp() - interval '1 second' WHERE id = $1", first.ID) + _, err = fixture.Store.pool.Exec( + ctx, + "UPDATE event_webhook_deliveries SET next_attempt_at = statement_timestamp() - interval '1 second' WHERE id = $1", + first.ID, + ) require.NoError(t, err) third, err := fixture.Store.Execution().ClaimEventWebhookDelivery(ctx, 128) require.NoError(t, err) require.Equal(t, first.ID, third.ID) require.Equal(t, first.InteractionUpdate, third.InteractionUpdate) require.Equal(t, int32(3), third.AttemptCount) - _, err = fixture.Store.pool.Exec(ctx, + _, err = fixture.Store.pool.Exec( + ctx, "UPDATE event_webhook_deliveries SET created_at = statement_timestamp() - interval '9 minutes 30 seconds' WHERE id = $1", - third.ID) + third.ID, + ) require.NoError(t, err) - retry, err = fixture.Store.Execution().RetryEventWebhookDelivery(ctx, third.ID, third.ClaimToken, time.Minute) + retry, err = fixture.Store.Execution(). + RetryEventWebhookDelivery(ctx, third.ID, third.ClaimToken, time.Minute) require.NoError(t, err) require.True(t, retry.GaveUp) require.NoError(t, fixture.Store.pool.QueryRow(ctx, @@ -379,14 +452,20 @@ func TestEventWebhookClaimsRecoverAndExpire(t *testing.T) { SET claim_token = $2, claim_expires_at = statement_timestamp() + interval '30 seconds' WHERE id = $1`, third.ID, third.ClaimToken) require.NoError(t, err) - _, err = fixture.Store.pool.Exec(ctx, - "UPDATE event_webhook_deliveries SET created_at = statement_timestamp() - interval '11 minutes' WHERE id = $1", first.ID) + _, err = fixture.Store.pool.Exec( + ctx, + "UPDATE event_webhook_deliveries SET created_at = statement_timestamp() - interval '11 minutes' WHERE id = $1", + first.ID, + ) require.NoError(t, err) deleted, err := fixture.Store.Execution().DeleteExpiredEventWebhookDeliveries(ctx, 500) require.NoError(t, err) require.Zero(t, deleted) - _, err = fixture.Store.pool.Exec(ctx, - "UPDATE event_webhook_deliveries SET claim_expires_at = statement_timestamp() - interval '1 second' WHERE id = $1", first.ID) + _, err = fixture.Store.pool.Exec( + ctx, + "UPDATE event_webhook_deliveries SET claim_expires_at = statement_timestamp() - interval '1 second' WHERE id = $1", + first.ID, + ) require.NoError(t, err) _, err = fixture.Store.Execution().ClaimEventWebhookDelivery(ctx, 128) require.ErrorIs(t, err, storeerr.ErrNotFound) @@ -410,11 +489,22 @@ func TestEventWebhookClaimsShareOrganizationLimits(t *testing.T) { require.NoError(t, err) otherOrgID, otherProjectID := uuid.New(), uuid.New() - _, err = pool.Exec(ctx, + _, err = pool.Exec( + ctx, "INSERT INTO orgs(id, name, created_at, updated_at) VALUES ($1, 'Other org', statement_timestamp(), statement_timestamp())", - otherOrgID) + otherOrgID, + ) require.NoError(t, err) - storagefixture.InsertProject(t, ctx, pool, otherOrgID, otherProjectID, "Other project", "other-project", fixture.Now) + storagefixture.InsertProject( + t, + ctx, + pool, + otherOrgID, + otherProjectID, + "Other project", + "other-project", + fixture.Now, + ) _, err = fixture.Store.Identity().AddOrgMembership(ctx, identitystore.AddOrgMembershipInput{ OrgID: otherOrgID, UserID: fixture.UserID, Role: "admin", }) @@ -444,10 +534,11 @@ event_webhook: pending.AddToolCallUpdate(fixture.AgentID, uuid.New(), "ready", nil) } pending.AddToolCallUpdate(otherAgent.ID, uuid.New(), "ready", nil) - tx, err := pool.Begin(ctx) + unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) - defer func() { _ = tx.Rollback(ctx) }() - require.NoError(t, fixture.Store.Execution().IntegrationCommitTxWithNotifications(ctx, tx, pending, "org limits test")) + defer func() { _ = unit.Rollback(ctx) }() + unit.Notifications().Restore(pending) + require.NoError(t, unit.Commit(ctx, "org limits test")) workers := []*executionstore.Store{fixture.Store.Execution(), newIntegrationStore(pool).Execution()} claimed := make([]executionstore.EventWebhookDelivery, 0, perOrgLimit) @@ -475,9 +566,11 @@ event_webhook: _, err = workers[0].ClaimEventWebhookDelivery(ctx, perOrgLimit) require.ErrorIs(t, err, storeerr.ErrNotFound) - _, err = pool.Exec(ctx, + _, err = pool.Exec( + ctx, "UPDATE event_webhook_deliveries SET claim_expires_at = statement_timestamp() - interval '1 second' WHERE id = $1", - claimed[2].ID) + claimed[2].ID, + ) require.NoError(t, err) recovered, err := workers[1].ClaimEventWebhookDelivery(ctx, perOrgLimit) require.NoError(t, err) @@ -525,10 +618,11 @@ func TestEventWebhookCleanupIsBoundedAndPreservesUnexpiredDeliveries(t *testing. toolIDs[i] = uuid.New() pending.AddToolCallUpdate(fixture.AgentID, toolIDs[i], "ready", nil) } - tx, err := fixture.Store.pool.Begin(ctx) + unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) - defer func() { _ = tx.Rollback(ctx) }() - require.NoError(t, fixture.Store.Execution().IntegrationCommitTxWithNotifications(ctx, tx, pending, "cleanup test")) + defer func() { _ = unit.Rollback(ctx) }() + unit.Notifications().Restore(pending) + require.NoError(t, unit.Commit(ctx, "cleanup test")) for i, toolID := range toolIDs { _, err := fixture.Store.pool.Exec(ctx, ` UPDATE event_webhook_deliveries @@ -542,8 +636,11 @@ func TestEventWebhookCleanupIsBoundedAndPreservesUnexpiredDeliveries(t *testing. require.EqualValues(t, expected, deleted) } var remaining int - require.NoError(t, fixture.Store.pool.QueryRow(ctx, - "SELECT count(*) FROM event_webhook_deliveries WHERE agent_id = $1", fixture.AgentID).Scan(&remaining)) + require.NoError(t, fixture.Store.pool.QueryRow( + ctx, + "SELECT count(*) FROM event_webhook_deliveries WHERE agent_id = $1", + fixture.AgentID, + ).Scan(&remaining)) require.Equal(t, 3, remaining) } @@ -565,10 +662,11 @@ func TestEventWebhookDeletedOwnersHaveNoDeliveryTarget(t *testing.T) { require.NotEmpty(t, target.URL) pending := notifications.NewTxNotifications() pending.AddToolCallUpdate(fixture.AgentID, uuid.New(), "ready", nil) - tx, err := fixture.Store.pool.Begin(ctx) + unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) - defer func() { _ = tx.Rollback(ctx) }() - require.NoError(t, fixture.Store.Execution().IntegrationCommitTxWithNotifications(ctx, tx, pending, "deleted owner")) + defer func() { _ = unit.Rollback(ctx) }() + unit.Notifications().Restore(pending) + require.NoError(t, unit.Commit(ctx, "deleted owner")) _, err = fixture.Store.pool.Exec(ctx, owner.query, owner.id) require.NoError(t, err) delivery, err := fixture.Store.Execution().ClaimEventWebhookDelivery(ctx, 128) @@ -577,7 +675,10 @@ func TestEventWebhookDeletedOwnersHaveNoDeliveryTarget(t *testing.T) { target, err = fixture.Store.Execution().GetAgentEventWebhookTarget(ctx, delivery.AgentID) require.NoError(t, err) require.Empty(t, target.URL) - require.NoError(t, fixture.Store.Execution().CompleteEventWebhookDelivery(ctx, delivery.ID, delivery.ClaimToken)) + require.NoError( + t, + fixture.Store.Execution().CompleteEventWebhookDelivery(ctx, delivery.ID, delivery.ClaimToken), + ) }) } } diff --git a/internal/storage/executionstore/event_webhooks_store.go b/internal/storage/executionstore/event_webhooks_store.go index fef12d995..41507a249 100644 --- a/internal/storage/executionstore/event_webhooks_store.go +++ b/internal/storage/executionstore/event_webhooks_store.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "slices" "time" "github.com/google/uuid" @@ -57,81 +56,6 @@ func (s *Store) GetAgentEventForWebhook( return agentEventReadRecordFromSQLC(rows[0]), nil } -func enqueueEventWebhooksTx(ctx context.Context, tx pgx.Tx, pending *notifications.TxNotifications) error { - if pending == nil || (len(pending.AgentEvents()) == 0 && len(pending.ToolCallUpdates()) == 0) { - return nil - } - q := dbsqlc.New(tx) - type webhookTarget struct { - orgID uuid.UUID - events []string - } - webhookTargets := make(map[uuid.UUID]webhookTarget) - for agentID := range pending.AgentEvents() { - webhookTargets[agentID] = webhookTarget{} - } - for _, update := range pending.ToolCallUpdates() { - webhookTargets[update.AgentID] = webhookTarget{} - } - for agentID := range webhookTargets { - target, err := q.GetAgentEventWebhookTarget(ctx, dbsqlc.GetAgentEventWebhookTargetParams{ID: agentID}) - if errors.Is(err, pgx.ErrNoRows) { - delete(webhookTargets, agentID) - continue - } - if err != nil { - return fmt.Errorf("load event webhook target: %w", err) - } - if target.Url == "" { - delete(webhookTargets, agentID) - continue - } - var events []string - if err := json.Unmarshal(target.Events, &events); err != nil { - return fmt.Errorf("decode event webhook events: %w", err) - } - webhookTargets[agentID] = webhookTarget{orgID: target.OrgID, events: events} - } - for agentID, events := range pending.AgentEvents() { - target, enabled := webhookTargets[agentID] - if !enabled { - continue - } - for _, event := range events { - if !slices.Contains(target.events, event.Kind) { - continue - } - if err := q.EnqueueEventWebhookDelivery(ctx, dbsqlc.EnqueueEventWebhookDeliveryParams{ - OrgID: target.orgID, AgentID: agentID, EventSequence: &event.Sequence, - }); err != nil { - return fmt.Errorf("enqueue event webhook: %w", err) - } - } - } - for _, update := range pending.ToolCallUpdates() { - target, enabled := webhookTargets[update.AgentID] - if !enabled || !slices.Contains(target.events, "tool_call_update") { - continue - } - var interactionUpdate json.RawMessage - if update.InteractionUpdate != nil { - var err error - interactionUpdate, err = json.Marshal(update.InteractionUpdate) - if err != nil { - return fmt.Errorf("encode webhook interaction update: %w", err) - } - } - if err := q.EnqueueEventWebhookDelivery(ctx, dbsqlc.EnqueueEventWebhookDeliveryParams{ - OrgID: target.orgID, AgentID: update.AgentID, - ToolCallID: &update.ToolCallID, ToolState: &update.State, - InteractionUpdate: sqlcRawMessageFromEmpty(interactionUpdate), - }); err != nil { - return fmt.Errorf("enqueue tool webhook: %w", err) - } - } - return nil -} - type EventWebhookDelivery struct { OrgID uuid.UUID ID uuid.UUID diff --git a/internal/storage/executionstore/execution_content.go b/internal/storage/executionstore/execution_content.go new file mode 100644 index 000000000..0a8a39e02 --- /dev/null +++ b/internal/storage/executionstore/execution_content.go @@ -0,0 +1,26 @@ +package executionstore + +import ( + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" +) + +func executionContent(blocks []CreateContentBlockInput) ([]agentexecution.Content, error) { + result := make([]agentexecution.Content, len(blocks)) + for i, block := range blocks { + metadata, err := block.Metadata.JSON() + if err != nil { + return nil, err + } + result[i] = agentexecution.Content{ + Ordinal: block.Ordinal, + Kind: string(block.BlockKind), + Text: block.TextContent, + Data: block.StructuredData, + ArtifactID: block.ArtifactID, + ToolID: block.ToolCallID, + Exclude: block.ExcludeFromModelContext, + Metadata: metadata, + } + } + return result, nil +} diff --git a/internal/storage/executionstore/execution_models.go b/internal/storage/executionstore/execution_models.go new file mode 100644 index 000000000..d39329ad8 --- /dev/null +++ b/internal/storage/executionstore/execution_models.go @@ -0,0 +1,65 @@ +package executionstore + +import ( + "context" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" +) + +func prepareExecutionModel( + ctx context.Context, + unit *agentexecution.Unit, + input PrepareNormalModelCallInput, +) (PreparedNormalModelCall, error) { + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return PreparedNormalModelCall{}, err + } + prepared, err := h.PrepareNormal(ctx, agentexecution.PrepareNormalInput{ + RuntimeLockID: input.RuntimeLockID, + OpeningInputIDs: input.OpeningInputIDs, + SourceContextID: input.SourceModelCallContextID, + SourceOutputID: input.SourceModelOutputID, + }) + if err != nil { + return PreparedNormalModelCall{}, err + } + return shapePreparedModel(ctx, unit, input.ProjectID, input.AgentID, prepared) +} + +func executionFailure(input RecordRecoverableModelCallFailureInput) agentexecution.ModelFailure { + return agentexecution.ModelFailure{ + ContextID: input.ModelCallContextID, + RuntimeLockID: input.RuntimeLockID, + Recovery: agentexecution.RecoveryKind(input.RecoveryKind), + RetryDelay: input.RetryDelay, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, APIVariant: input.APIVariant, + RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, Usage: input.Usage, + Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}, + } +} + +func executionClaim( + ctx context.Context, + unit *agentexecution.Unit, + projectID, agentID uuid.UUID, + prepared agentexecution.PreparedModel, +) (ModelCallClaim, error) { + if prepared.Context.ID == uuid.Nil { + return ModelCallClaim{}, nil + } + record, err := loadModelCallContextByID( + ctx, + dbsqlc.New(unit.DB()), + projectID, + agentID, + prepared.Context.ID, + ) + return ModelCallClaim{Context: record, Claimed: prepared.Claimed, Created: prepared.Created}, err +} diff --git a/internal/storage/executionstore/execution_mutation.go b/internal/storage/executionstore/execution_mutation.go new file mode 100644 index 000000000..90d839600 --- /dev/null +++ b/internal/storage/executionstore/execution_mutation.go @@ -0,0 +1,30 @@ +package executionstore + +import ( + "context" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" +) + +func beginExecution( + ctx context.Context, + s *Store, + projectID, agentID uuid.UUID, +) (*agentexecution.Unit, *agentexecution.Handle, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return nil, nil, err + } + if err = lockAgentWithParentTx(ctx, unit, dbsqlc.New(unit.DB()), projectID, agentID); err != nil { + _ = unit.Rollback(ctx) + return nil, nil, err + } + h, err := unit.Handle(projectID, agentID) + if err != nil { + _ = unit.Rollback(ctx) + return nil, nil, err + } + return unit, h, nil +} diff --git a/internal/storage/executionstore/execution_unit_lock_order_integration_test.go b/internal/storage/executionstore/execution_unit_lock_order_integration_test.go new file mode 100644 index 000000000..91d900271 --- /dev/null +++ b/internal/storage/executionstore/execution_unit_lock_order_integration_test.go @@ -0,0 +1,329 @@ +//go:build integration + +package executionstore_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + "github.com/stretchr/testify/require" +) + +func executionLockOrderAgents(t *testing.T, parentFirst bool) (processDaemonFixture, uuid.UUID, uuid.UUID) { + t.Helper() + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "execution-unit-locks") + parent := uuid.MustParse("00000000-0000-0000-0000-000000000001") + child := uuid.MustParse("ffffffff-ffff-ffff-ffff-ffffffffffff") + if !parentFirst { + parent, child = child, parent + } + original, err := f.Store.Execution().GetAgentInProject(ctx, testProjectID, f.AgentID) + require.NoError(t, err) + unit, err := f.Store.Execution().IntegrationBeginUnit(ctx) + require.NoError(t, err) + defer func() { _ = unit.Rollback(ctx) }() + _, err = unit.CreateAgent( + ctx, + agentexecution.CreateAgentInput{ + ID: parent, + ProjectID: testProjectID, + ConfigID: original.CurrentConfigID, + Name: "parent", + }, + ) + require.NoError(t, err) + _, err = unit.CreateAgent( + ctx, + agentexecution.CreateAgentInput{ + ID: child, + ProjectID: testProjectID, + ConfigID: original.CurrentConfigID, + ParentID: parent, + SubagentKey: "fork", + Name: "child", + }, + ) + require.NoError(t, err) + require.NoError(t, unit.Commit(ctx, "seed family")) + return f, parent, child +} + +func assertExecutionAgentOrder( + t *testing.T, f processDaemonFixture, first, second uuid.UUID, run func(context.Context) error, +) { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), 20*time.Second) + defer cancel() + blocker := integrationdb.BeginTx(t, ctx, f.Store.pool) + _, err := blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, first) + require.NoError(t, err) + done := integrationdb.RunAsyncError(func() error { return run(ctx) }) + integrationdb.WaitForNamedLockWaiters(t, ctx, f.Store.pool, "LockAgentInProject", 1) + _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, second) + require.NoError(t, err, "the later agent must remain free while waiting for the first") + require.NoError(t, blocker.Commit(ctx)) + require.ErrorIs( + t, + integrationdb.Await(t, done, "ordered execution mutation"), + storeerr.ErrRuntimeLockInactive, + ) +} + +func TestModelMutationPlansParentBeforeRuntime(t *testing.T) { + for _, parentFirst := range []bool{true, false} { + t.Run(fmt.Sprintf("parent_first_%t", parentFirst), func(t *testing.T) { + f, parent, child := executionLockOrderAgents(t, parentFirst) + first, second := parent, child + if !parentFirst { + first, second = child, parent + } + runtime, source := uuid.New(), uuid.New() + for _, operation := range []string{ + "claim normal", "claim compaction", "retry", "prepare", + "begin compaction", "replace compaction", "terminal compaction", + } { + t.Run(operation, func(t *testing.T) { + assertExecutionAgentOrder(t, f, first, second, func(ctx context.Context) error { + s := f.Store.Execution() + switch operation { + case "claim normal": + _, err := s.PrepareNormalModelCall( + ctx, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, AgentID: child, RuntimeLockID: runtime, + OpeningInputIDs: []uuid.UUID{source}}, + ) + return err + case "claim compaction": + _, err := s.ClaimCompactionModelCall( + ctx, + executionstore.ClaimCompactionModelCallInput{ + ProjectID: testProjectID, AgentID: child, RuntimeLockID: runtime, ParentContextID: source, + InputEventSequence: 1, SourceEventSequenceEnd: 1}, + ) + return err + case "retry": + _, err := s.ClaimNextModelCallContext( + ctx, + executionstore.ClaimNextModelCallContextInput{ + ProjectID: testProjectID, + AgentID: child, + RuntimeLockID: runtime, + PredecessorModelCallContextID: source}, + ) + return err + case "prepare": + _, err := s.PrepareNormalModelCall( + ctx, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: child, + RuntimeLockID: runtime, + OpeningInputIDs: []uuid.UUID{source}}, + ) + return err + case "begin compaction": + _, err := s.RecordModelCallFailureAndClaimCompaction(ctx, + executionstore.RecordModelCallFailureAndClaimCompactionInput{ + ParentContextID: source, SourceEventSequenceEnd: 1, + Failure: executionstore.RecordRecoverableModelCallFailureInput{ + ProjectID: testProjectID, AgentID: child, RuntimeLockID: runtime, ModelCallContextID: source, + RecoveryKind: executionstore.ModelCallRecoveryCompact, + ErrorKind: "context_length", ErrorMessage: "too large"}}) + return err + case "replace compaction": + _, err := s.ReplaceCompactionSource( + ctx, + executionstore.ReplaceCompactionSourceInput{ + ProjectID: testProjectID, AgentID: child, RuntimeLockID: runtime, ModelCallContextID: source, + NextSourceEventSequenceEnd: 1, ErrorKind: "context_length", ErrorMessage: "too large"}, + ) + return err + default: + return s.RecordTerminalCompactionFailure( + ctx, + executionstore.RecordTerminalCompactionFailureInput{ + ProjectID: testProjectID, AgentID: child, RuntimeLockID: runtime, ModelCallContextID: source, + ErrorKind: "context_length", ErrorMessage: "too large"}, + ) + } + }) + }) + } + }) + } +} + +func TestStopSubagentPlansCallerAndTreeBeforeMutation(t *testing.T) { + for _, parentFirst := range []bool{true, false} { + for _, archive := range []bool{false, true} { + t.Run(fmt.Sprintf("parent_first_%t_archive_%t", parentFirst, archive), func(t *testing.T) { + f, parent, child := executionLockOrderAgents(t, parentFirst) + first, second := parent, child + if !parentFirst { + first, second = child, parent + } + assertExecutionAgentOrder(t, f, first, second, func(ctx context.Context) error { + _, err := f.Store.Execution().ExecuteToolCall(ctx, executionstore.ExecuteToolCallInput{ + ProjectID: testProjectID, AgentID: parent, ToolCallID: uuid.New(), RuntimeLockID: uuid.New(), + }, func(*executionstore.ToolCallReader) (executionstore.ToolCallCommand, error) { + return executionstore.StopSubagentForToolCall(executionstore.StopSubagentInput{ + TargetAgentID: child, Archive: archive, + }, executionstore.ToolCallCompletionInput{}), nil + }) + return err + }) + agent, err := f.Store.Execution().GetAgentInProject(t.Context(), testProjectID, child) + require.NoError(t, err) + require.Equal(t, executionstore.AgentStateActive, agent.State) + }) + } + } +} + +func TestSubagentLaunchLocksMachinesBeforeParentAndInsert(t *testing.T) { + t.Parallel() + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "subagent-machine-order") + parent, err := f.Store.Execution().GetAgentInProject(ctx, testProjectID, f.AgentID) + require.NoError(t, err) + blocker := integrationdb.BeginTx(t, ctx, f.Store.pool) + _, err = blocker.Exec(ctx, `SELECT id FROM machines WHERE id=$1 FOR UPDATE`, f.MachineID) + require.NoError(t, err) + done := integrationdb.RunAsync(func() (executionstore.LaunchAgentResult, error) { + return spawnSubagentForTest( + t, + ctx, + f.Store, + parent, + parent.CurrentConfigID, + "helper", + "machine-order-child", + nil, + ) + }) + integrationdb.WaitForNamedLockWaiters(t, ctx, f.Store.pool, "LockMachineForLifecycle", 1) + _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, parent.ID) + require.NoError(t, err) + var children int + require.NoError(t, blocker.QueryRow(ctx, + `SELECT count(*) FROM agents WHERE parent_agent_id=$1`, parent.ID).Scan(&children)) + require.Zero(t, children) + require.NoError(t, blocker.Commit(ctx)) + outcome := integrationdb.Await(t, done, "subagent launch") + require.NoError(t, outcome.Err) + require.Equal(t, parent.ID, outcome.Value.Agent.ParentAgentID) + require.Len(t, outcome.Value.MachineBindings, 1) +} + +func TestSubagentLaunchTakesQuotaBeforeParentAndMachines(t *testing.T) { + t.Parallel() + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "subagent-quota-order") + parent, err := f.Store.Execution().GetAgentInProject(ctx, testProjectID, f.AgentID) + require.NoError(t, err) + blocker := integrationdb.BeginTx(t, ctx, f.Store.pool) + require.NoError(t, dbsqlc.New(blocker).LockResourceCreation(ctx, dbsqlc.LockResourceCreationParams{ + ResourceKind: "agents", Scope: testProjectID.String(), + })) + done := integrationdb.RunAsync(func() (executionstore.LaunchAgentResult, error) { + return spawnSubagentForTest( + t, + ctx, + f.Store, + parent, + parent.CurrentConfigID, + "helper", + "quota-order-child", + nil, + ) + }) + integrationdb.WaitForNamedLockWaiters(t, ctx, f.Store.pool, "LockResourceCreation", 1) + _, err = blocker.Exec(ctx, `SELECT id FROM machines WHERE id=$1 FOR UPDATE NOWAIT`, f.MachineID) + require.NoError(t, err) + _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, parent.ID) + require.NoError(t, err) + require.NoError(t, blocker.Commit(ctx)) + require.NoError(t, integrationdb.Await(t, done, "subagent launch after quota").Err) +} + +func TestPreparationFencesLeaseAfterParentWait(t *testing.T) { + t.Parallel() + ctx := t.Context() + f, parent, child := executionLockOrderAgents(t, true) + runtime, err := f.Store.Execution().AcquireAgentRuntimeLock( + ctx, testProjectID, child, testWorkerProcessID, time.Minute, + ) + require.NoError(t, err) + blocker := integrationdb.BeginTx(t, ctx, f.Store.pool) + _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, parent) + require.NoError(t, err) + done := integrationdb.RunAsyncError(func() error { + _, err := f.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: child, + RuntimeLockID: runtime.ID, + OpeningInputIDs: []uuid.UUID{uuid.New()}, + }) + return err + }) + integrationdb.WaitForNamedLockWaiters(t, ctx, f.Store.pool, "LockAgentInProject", 1) + expireAgentRuntimeLockForTest(t, ctx, f.Store, runtime.ID) + require.NoError(t, blocker.Commit(ctx)) + require.ErrorIs( + t, + integrationdb.Await(t, done, "preparation after lease expiry"), + storeerr.ErrRuntimeLockInactive, + ) +} + +func TestSubagentLaunchReplansAfterMachineTeardown(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "subagent-machine-teardown") + parent, err := f.Store.Execution().GetAgentInProject(ctx, testProjectID, f.AgentID) + require.NoError(t, err) + teardown, err := f.Store.Execution().IntegrationBeginUnit(ctx) + require.NoError(t, err) + defer func() { _ = teardown.Rollback(ctx) }() + _, err = dbsqlc.New(teardown.DB()).LockMachineForLifecycle(ctx, dbsqlc.LockMachineForLifecycleParams{ + OrgID: testOrgID, ID: f.MachineID, + }) + require.NoError(t, err) + done := integrationdb.RunAsync(func() (executionstore.LaunchAgentResult, error) { + return spawnSubagentForTest( + t, + ctx, + f.Store, + parent, + parent.CurrentConfigID, + "helper", + "teardown-child", + nil, + ) + }) + integrationdb.WaitForNamedLockWaiters(t, ctx, f.Store.pool, "LockMachineForLifecycle", 1) + _, err = teardown.DB().Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, parent.ID) + require.NoError(t, err) + _, err = f.Store.Execution().DeleteMachineInUnit(ctx, teardown, executionstore.DeleteMachineInput{ + OrgID: testOrgID, MachineID: f.MachineID, + }) + require.NoError(t, err) + require.NoError(t, teardown.Commit(ctx, "teardown before launch")) + result := integrationdb.Await(t, done, "launch after machine teardown") + require.NoError(t, result.Err) + require.True(t, result.Value.Created) + require.Empty(t, result.Value.MachineBindings) + var children int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM agents WHERE parent_agent_id=$1`, parent.ID).Scan(&children)) + require.Equal(t, 1, children) +} diff --git a/internal/storage/executionstore/execution_unit_webhooks_integration_test.go b/internal/storage/executionstore/execution_unit_webhooks_integration_test.go new file mode 100644 index 000000000..2efdfaba1 --- /dev/null +++ b/internal/storage/executionstore/execution_unit_webhooks_integration_test.go @@ -0,0 +1,84 @@ +//go:build integration + +package executionstore_test + +import ( + "context" + "fmt" + "testing" + + "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/stretchr/testify/require" +) + +func TestLifecycleUnitEnqueuesWebhooksAtomically(t *testing.T) { + for _, scope := range []string{"project", "organization"} { + for _, fail := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/fail_%t", scope, fail), func(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "lifecycle-webhooks") + toolID := createToolCallForProcessTestWithPermission(t, ctx, f, "lifecycle-question", "ask_question", true) + enableEventWebhook(t, f, []string{"agent_input", "tool_result", "tool_call_update"}) + question := createQuestionInteractionForTest(t, ctx, f, toolID) + publisher := &recordingPostCommitPublisher{} + f.Store = newIntegrationStore(f.Store.pool, storage.WithPostCommitPublisher(publisher)) + _, err := f.Store.pool.Exec(ctx, `DELETE FROM event_webhook_deliveries WHERE agent_id=$1`, f.AgentID) + require.NoError(t, err) + if fail { + _, err := f.Store.pool.Exec(ctx, ` +CREATE FUNCTION reject_unit_webhook() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN RAISE EXCEPTION 'rejected webhook'; END $$; +CREATE TRIGGER reject_unit_webhook BEFORE INSERT ON event_webhook_deliveries +FOR EACH ROW EXECUTE FUNCTION reject_unit_webhook();`) + require.NoError(t, err) + } + remove := func(ctx context.Context) error { + if scope == "project" { + _, err := f.Store.Organizations().DeleteProject(ctx, testOrgID, testProjectID, userPrincipal(f.UserID)) + return err + } + _, err := f.Store.Organizations().DeleteOrganization(ctx, testOrgID, userPrincipal(f.UserID)) + return err + } + err = remove(ctx) + if fail { + require.ErrorContains(t, err, "EnqueueEventWebhookDelivery") + } else { + require.NoError(t, err) + } + var state string + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT state FROM agents WHERE id=$1`, f.AgentID).Scan(&state)) + var count int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM event_webhook_deliveries WHERE agent_id=$1`, f.AgentID).Scan(&count)) + if fail { + require.Equal(t, "active", state) + require.Zero(t, count) + require.Empty(t, publisher.intents) + var interactionState string + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT state FROM agent_interactions WHERE id=$1`, question.ID).Scan(&interactionState)) + require.Equal(t, string(executionstore.AgentInteractionStateOpen), interactionState) + return + } + require.Equal(t, "archived", state) + require.Positive(t, count) + var completed int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM event_webhook_deliveries WHERE agent_id=$1 AND tool_call_id=$2 AND tool_state='completed'`, f.AgentID, toolID).Scan(&completed)) + require.Equal(t, 1, completed) + var published bool + for _, intent := range publisher.intents { + if update, ok := intent.(notifications.ToolCallUpdatedCommitted); ok && + update.ToolCallID == toolID && update.State == "completed" { + published = true + } + } + require.True(t, published) + }) + } + } +} diff --git a/internal/storage/executionstore/execution_work.go b/internal/storage/executionstore/execution_work.go new file mode 100644 index 000000000..b332e5991 --- /dev/null +++ b/internal/storage/executionstore/execution_work.go @@ -0,0 +1,153 @@ +package executionstore + +import ( + "context" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/events" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" +) + +func shapeOwnedWork( + ctx context.Context, + unit *agentexecution.Unit, + projectID, agentID uuid.UUID, + work agentexecution.OwnedWork, +) (ClaimedAgentWork, error) { + result := ClaimedAgentWork{ProjectID: projectID, AgentID: agentID} + if work.Released || work.Selection.Work == agentexecution.WorkNone { + return result, nil + } + q := dbsqlc.New(unit.DB()) + agent, err := q.GetAgentInProject(ctx, dbsqlc.GetAgentInProjectParams{ProjectID: projectID, ID: agentID}) + if err != nil { + return result, err + } + result.OrgID = agent.OrgID + runtime, err := q.GetAgentRuntimeLockForRelease( + ctx, + dbsqlc.GetAgentRuntimeLockForReleaseParams{ + ProjectID: projectID, + AgentID: agentID, + ID: work.RuntimeID, + }, + ) + if err != nil { + return result, err + } + result.RuntimeLock = agentRuntimeLockRecordFromSQLC(runtime) + if selected := work.Selection.Tool; selected != nil { + result.Kind = AgentWorkTool + result.Tool = ClaimedToolWork{TurnID: selected.TurnID, ModelCallContextID: selected.SourceContextID, + ModelOutputID: selected.OutputID, SourceEventID: selected.SourceEventID} + result.Tool.Prepared, err = prepareToolWorkTx(ctx, q, projectID, agentID, result.Tool) + return result, err + } + selected := work.Selection.Model + result.Kind = AgentWorkModel + result.Model = ClaimedModelWork{Kind: ModelWorkKind(selected.Kind), TurnID: selected.TurnID, + InputIDs: selected.Opening.InputIDs, OpeningEventSequence: selected.Opening.EventSequence} + if selected.Kind == agentexecution.ModelResume { + result.Model.ModelCallContextID = selected.SourceContextID + } + if selected.Kind == agentexecution.ModelContinue { + result.Model.SourceModelCallContextID = selected.SourceContextID + result.Model.SourceModelOutputID = selected.SourceOutputID + } + if work.Model.Context.ID != uuid.Nil { + prepared, shapeErr := shapePreparedModel(ctx, unit, projectID, agentID, work.Model) + if shapeErr != nil { + return result, shapeErr + } + result.Model.Prepared = &prepared + } + if work.Admission.TurnID != uuid.Nil { + result.Model.AdmittedInputTurn, err = shapeAdmission(ctx, q, projectID, agentID, work.Admission) + if err != nil { + return result, err + } + } + + return result, nil +} + +func shapePreparedModel( + ctx context.Context, + unit *agentexecution.Unit, + projectID, agentID uuid.UUID, + prepared agentexecution.PreparedModel, +) (PreparedNormalModelCall, error) { + q := dbsqlc.New(unit.DB()) + record, err := loadModelCallContextByID(ctx, q, projectID, agentID, prepared.Context.ID) + if err != nil { + return PreparedNormalModelCall{}, err + } + config, err := q.GetAgentConfig( + ctx, + dbsqlc.GetAgentConfigParams{ProjectID: projectID, ID: record.AgentConfigID}, + ) + if err != nil { + return PreparedNormalModelCall{}, err + } + result := PreparedNormalModelCall{ + Claim: ModelCallClaim{Context: record, Created: prepared.Created, Claimed: prepared.Claimed}, + Snapshot: AgentConfigSnapshotRecord{ + AgentConfig: agentConfigRecordFromSQLC(config), + InputEventSequence: record.InputEventSequence, + }, + } + if prepared.Claimed { + result.ContextData, err = loadModelContextDataTx(ctx, q, record) + } + return result, err +} + +func shapeAdmission( + ctx context.Context, + q *dbsqlc.Queries, + projectID, agentID uuid.UUID, + admission agentexecution.InputAdmission, +) (AdmittedAgentInputTurn, error) { + var result AdmittedAgentInputTurn + if len(admission.Inputs) == 0 { + return result, nil + } + turn, err := q.GetAgentTurn( + ctx, + dbsqlc.GetAgentTurnParams{ProjectID: projectID, AgentID: agentID, ID: admission.TurnID}, + ) + if err != nil { + return result, err + } + result.Turn = AgentTurnRecord{ + ID: turn.ID, + ProjectID: projectID, + AgentID: agentID, + TurnSequence: turn.TurnSequence, + LatestEventID: turn.LatestEventID, + LatestSemanticEventID: turn.LatestSemanticEventID, + } + for _, admitted := range admission.Inputs { + row, err := q.GetAgentInput( + ctx, + dbsqlc.GetAgentInputParams{ProjectID: projectID, AgentID: agentID, ID: admitted.ID}, + ) + if err != nil { + return result, err + } + result.Inputs = append(result.Inputs, agentInputRecordFromGetSQLC(row)) + result.Events = append( + result.Events, + events.Event{ + ID: admitted.Event.ID, + AgentID: agentID, + Sequence: admitted.Event.Sequence, + Kind: events.KindAgentInput, + At: admitted.Event.Time, + IdempotencyKey: "agent_input:" + admitted.ID.String(), + }, + ) + } + return result, nil +} diff --git a/internal/storage/executionstore/export_integration_test.go b/internal/storage/executionstore/export_integration_test.go index 8318c458e..ff005be51 100644 --- a/internal/storage/executionstore/export_integration_test.go +++ b/internal/storage/executionstore/export_integration_test.go @@ -6,7 +6,6 @@ import ( "context" "encoding/json" "errors" - "fmt" "time" "github.com/google/uuid" @@ -15,7 +14,11 @@ import ( "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" + "github.com/omnara-ai/omnara/internal/storage/storeerr" ) type ModelWorkSeed struct { @@ -27,132 +30,116 @@ type ModelWorkSeed struct { OpeningEventSequence int64 } -func (s *Store) IntegrationCommitTxWithNotifications( - ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, - operation string, -) error { - return s.commitTxWithNotifications(ctx, tx, txNotifications, operation) +func (s *Store) IntegrationBeginUnit(ctx context.Context) (*agentexecution.Unit, error) { + return s.cell.Begin(ctx) } func (s *Store) AcquireAgentRuntimeLock( ctx context.Context, - projectID, agentID, workerProcessID uuid.UUID, - leaseDuration time.Duration, + projectID, agentID, workerID uuid.UUID, + lease time.Duration, ) (AgentRuntimeLockRecord, error) { - if projectID == uuid.Nil || agentID == uuid.Nil || workerProcessID == uuid.Nil { - return AgentRuntimeLockRecord{}, errors.New("project, agent, and worker process ids are required") + if err := validateAgentRuntimeLockLeaseDuration(lease); err != nil { + return AgentRuntimeLockRecord{}, err } - if err := validateAgentRuntimeLockLeaseDuration(leaseDuration); err != nil { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if err != nil { return AgentRuntimeLockRecord{}, err } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + defer func() { _ = unit.Rollback(ctx) }() + snapshot, err := h.LoadExecution(ctx) if err != nil { - return AgentRuntimeLockRecord{}, fmt.Errorf("begin acquire agent runtime lock: %w", err) + return AgentRuntimeLockRecord{}, err } - defer func() { _ = tx.Rollback(ctx) }() - record, _, err := acquireAgentRuntimeLockTx( - ctx, - dbsqlc.New(tx), - projectID, - agentID, - workerProcessID, - leaseDuration, - ) + if snapshot.View.State != agentexecution.AgentActive { + return AgentRuntimeLockRecord{}, storeerr.ErrAgentNotAdvanceable + } + var record AgentRuntimeLockRecord + err = unit.DB(). + QueryRow(ctx, + `INSERT INTO agent_runtime_locks(agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + VALUES($1,$2,statement_timestamp(),statement_timestamp(),statement_timestamp()+$3::bigint*interval '1 microsecond') + RETURNING id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at,cancel_requested_at`, + agentID, + workerID, + lease.Microseconds()). + Scan(&record.ID, + &record.AgentID, + &record.WorkerProcessID, + &record.StartedAt, + &record.RenewedAt, + &record.LeaseExpiresAt, + &record.CancelRequestedAt) if err != nil { return AgentRuntimeLockRecord{}, err } - if err := tx.Commit(ctx); err != nil { - return AgentRuntimeLockRecord{}, fmt.Errorf("commit acquire agent runtime lock: %w", err) + if _, err = h.Repair(ctx); err != nil { + return AgentRuntimeLockRecord{}, err } - return record, nil + return record, unit.Commit(ctx, "seed runtime") } -func (s *Store) MarkAgentWakeup( - ctx context.Context, - projectID, agentID uuid.UUID, - metadata []byte, -) error { - if projectID == uuid.Nil || agentID == uuid.Nil { - return errors.New("project id and agent id are required") - } - if metadata == nil { - metadata = []byte(`{}`) - } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) +func (s *Store) RepairExecutionForTest(ctx context.Context, projectID, agentID uuid.UUID) error { + unit, h, err := beginExecution(ctx, s, projectID, agentID) if err != nil { - return fmt.Errorf("begin mark agent wakeup: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: projectID, ID: agentID}, - ); err != nil { - return fmt.Errorf("lock agent for wakeup: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - Metadata: metadata, - }, - ); err != nil { - return fmt.Errorf("mark agent wakeup: %w", err) + return err } - if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("commit mark agent wakeup: %w", err) + defer func() { _ = unit.Rollback(ctx) }() + if _, err = h.Repair(ctx); err != nil { + return err } - return nil + return unit.Commit(ctx, "repair execution") } func (s *Store) DeleteAgentWakeup(ctx context.Context, projectID, agentID uuid.UUID) error { - if projectID == uuid.Nil || agentID == uuid.Nil { - return errors.New("project id and agent id are required") - } - if _, err := s.q.DeleteAgentWakeup( + _, err := s.pool.Exec( ctx, - dbsqlc.DeleteAgentWakeupParams{ProjectID: projectID, AgentID: agentID}, - ); err != nil { - return fmt.Errorf("delete agent wakeup: %w", err) - } - return nil + `DELETE FROM agent_wakeups w USING agents a WHERE a.id=w.agent_id AND a.project_id=$1 AND a.id=$2`, + projectID, + agentID, + ) + return err } func (s *Store) NextAgentModelWork( ctx context.Context, projectID, agentID uuid.UUID, ) (ModelWorkSeed, bool, error) { - if projectID == uuid.Nil || agentID == uuid.Nil { - return ModelWorkSeed{}, false, errors.New("project id and agent id are required") - } - row, err := s.q.NextAgentModelWork( - ctx, - dbsqlc.NextAgentModelWorkParams{ProjectID: projectID, AgentID: agentID}, - ) - if errors.Is(err, pgx.ErrNoRows) { - return ModelWorkSeed{}, false, nil + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if err != nil { + return ModelWorkSeed{}, false, err } + defer func() { _ = unit.Rollback(ctx) }() + snapshot, err := h.LoadExecution(ctx) if err != nil { - return ModelWorkSeed{}, false, fmt.Errorf("next agent model work: %w", err) + return ModelWorkSeed{}, false, err + } + work := snapshot.Selection.Model + if work == nil { + return ModelWorkSeed{}, false, nil } return ModelWorkSeed{ - Kind: ModelWorkKind(row.WorkKind), - ModelCallContextID: row.ModelCallContextID, - SourceModelOutputID: row.ModelOutputID, - TurnID: row.TurnID, - InputIDs: row.InputIds, - OpeningEventSequence: row.OpeningEventSequence, + Kind: ModelWorkKind(work.Kind), + ModelCallContextID: work.SourceContextID, + SourceModelOutputID: work.SourceOutputID, + TurnID: work.TurnID, + InputIDs: work.Opening.InputIDs, + OpeningEventSequence: work.Opening.EventSequence, }, true, nil } func (s *Store) GetToolCallResultAuthorityByToolCall( ctx context.Context, - projectID, agentID, toolCallID uuid.UUID, + projectID, agentID, toolID uuid.UUID, ) (ToolCallResultAuthorityRecord, bool, error) { - return getToolCallResultAuthorityByToolCallTx(ctx, s.pool, projectID, agentID, toolCallID) + row, err := s.q.GetToolCallResultByToolCall( + ctx, + dbsqlc.GetToolCallResultByToolCallParams{ProjectID: projectID, AgentID: agentID, ToolCallID: toolID}, + ) + if errors.Is(err, pgx.ErrNoRows) { + return ToolCallResultAuthorityRecord{}, false, nil + } + return toolCallResultAuthorityFromGetSQLC(row), err == nil, err } func (s *Store) RegisterDaemonRuntime( @@ -236,53 +223,20 @@ type IntegrationPoolMachineBindingInput struct { func IntegrationEnsureRuntimeLockActiveTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, projectID, agentID, runtimeID uuid.UUID, ) error { - return ensureRuntimeLockActiveTx(ctx, tx, projectID, agentID, runtimeID) -} - -func IntegrationSelectLockedSteeringAgentInputsForAdmissionTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) ([]AgentInputRecord, error) { - return selectLockedSteeringAgentInputsForAdmissionTx(ctx, qtx, projectID, agentID) -} - -func IntegrationSelectLockedQueuedAgentInputForAdmissionTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) ([]AgentInputRecord, error) { - return selectLockedQueuedAgentInputForAdmissionTx(ctx, qtx, projectID, agentID) -} - -func IntegrationAdmitLockedAgentInputsAndOpenTurnTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - input IntegrationAdmitAgentInputAndOpenTurnInput, - lockedInputs []AgentInputRecord, -) (AdmittedAgentInputTurn, error) { - return admitLockedAgentInputsAndOpenTurnTx( - ctx, - txNotifications, - tx, - qtx, - admitAgentInputAndOpenTurnInput(input), - lockedInputs, - ) -} - -func IntegrationModelCallOpeningInputSet( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, turnID uuid.UUID, - inputEventSequence int64, -) ([]uuid.UUID, int64, error) { - return modelCallOpeningInputSet(ctx, qtx, projectID, agentID, turnID, inputEventSequence) + if err := unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: projectID, AgentID: agentID}}, + agentexecution.RuntimeAuthority{AgentID: agentID, RuntimeLockID: runtimeID}, + ); err != nil { + return err + } + h, err := unit.Handle(projectID, agentID) + if err != nil { + return err + } + return h.FenceRuntime(ctx, runtimeID) } func IntegrationParseAgentInputContentBlocks( @@ -293,13 +247,12 @@ func IntegrationParseAgentInputContentBlocks( func IntegrationCreateAgentContentInputTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input CreateAgentContentInputInput, contentBlocks []CreateContentBlockInput, ) (IntegrationCreateAgentContentInputTxResult, error) { - result, err := createAgentContentInputTx(ctx, txNotifications, tx, qtx, input, contentBlocks) + result, err := createAgentContentInputTx(ctx, unit, qtx, input, contentBlocks) return IntegrationCreateAgentContentInputTxResult{ AgentInput: result.agentInput, ContentBlocks: result.contentBlocks, @@ -372,38 +325,54 @@ func (s *Store) IntegrationResolveLaunchMachineSourcesTx( func IntegrationValidateResponseEnvelopeForModelCallContext( ctx context.Context, - qtx *dbsqlc.Queries, + db dbsqlc.DBTX, envelope modelenvelope.ResponseEnvelope, - contextRow ModelCallContextRecord, + record ModelCallContextRecord, ) error { - return validateResponseEnvelopeForModelCallContext(ctx, qtx, envelope, contextRow) + var slug string + if err := db.QueryRow(ctx, + `SELECT provider_model_slug FROM configured_model_revisions WHERE id=$1`, + record.ConfiguredModelRevisionID).Scan(&slug); err != nil { + return err + } + if slug != envelope.RequestedProviderModelSlug { + return errors.New("requested model differs from captured revision") + } + return nil } func IntegrationRenewAgentRuntimeLockTx( ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, runtimeLockID uuid.UUID, - leaseDuration time.Duration, + unit *agentexecution.Unit, + projectID, agentID, runtimeID uuid.UUID, + lease time.Duration, ) (AgentRuntimeLockRenewal, error) { - return renewAgentRuntimeLockTx(ctx, qtx, projectID, agentID, runtimeLockID, leaseDuration) + renewal, err := unit.RenewRuntime( + ctx, + agentexecution.AgentRoute{CellID: unit.CellID(), ProjectID: projectID, AgentID: agentID}, + runtimeID, + lease, + ) + if err != nil { + return AgentRuntimeLockRenewal{}, err + } + row, err := dbsqlc.New(unit.DB()). + GetAgentRuntimeLockForRelease(ctx, + dbsqlc.GetAgentRuntimeLockForReleaseParams{ProjectID: projectID, + AgentID: agentID, + ID: runtimeID}) + return AgentRuntimeLockRenewal{ + RuntimeLock: agentRuntimeLockRecordFromSQLC(row), + LocalLeaseBudgetStartedAt: renewal.LocalStartedAt, + }, err } func IntegrationReapExpiredAgentRuntimeLockTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, projectID, agentID, runtimeLockID uuid.UUID, - retryBackoff func(int, string) time.Duration, ) (bool, error) { - return reapExpiredAgentRuntimeLockTx( - ctx, - txNotifications, - tx, - projectID, - agentID, - runtimeLockID, - retryBackoff, - ) + return reapExpiredAgentRuntimeLockTx(ctx, unit, projectID, agentID, runtimeLockID) } func IntegrationUpsertActorIdentityTx( @@ -416,10 +385,16 @@ func IntegrationUpsertActorIdentityTx( func IntegrationCompactionSourceStartTx( ctx context.Context, - qtx *dbsqlc.Queries, - contextRow ModelCallContextRecord, + db dbsqlc.DBTX, + row ModelCallContextRecord, ) (int64, error) { - return compactionSourceStartTx(ctx, qtx, contextRow) + var n int64 + err := db.QueryRow(ctx, + `SELECT coalesce(max(summarized_through_event_sequence),0)+1 FROM context_checkpoints WHERE agent_id=$1 AND summarized_through_event_sequence<$2`, + row.AgentID, + row.SourceEventSequenceEnd). + Scan(&n) + return n, err } func IntegrationCreateModelOutputAuthorityTx( @@ -427,7 +402,28 @@ func IntegrationCreateModelOutputAuthorityTx( db dbsqlc.DBTX, input CreateModelOutputAuthorityInput, ) (ModelOutputAuthorityRecord, error) { - return createModelOutputAuthorityTx(ctx, db, input) + var id uuid.UUID + err := db.QueryRow(ctx, + `INSERT INTO model_outputs(agent_id,model_call_context_id,served_provider_model_slug,stop_reason,provider_replay,created_at) + SELECT c.agent_id,c.id,$4,$5,$6,statement_timestamp() FROM model_call_contexts c WHERE + c.project_id=$1 AND c.agent_id=$2 AND c.id=$3 + RETURNING id`, + input.ProjectID, + input.AgentID, + input.ModelCallContextID, + input.ServedProviderModelSlug, + input.StopReason, + normalizedJSON(input.ProviderReplay)). + Scan(&id) + if err != nil { + return ModelOutputAuthorityRecord{}, err + } + row, err := dbsqlc.New(db). + GetModelOutputByModelContext(ctx, + dbsqlc.GetModelOutputByModelContextParams{ProjectID: input.ProjectID, + AgentID: input.AgentID, + ModelCallContextID: input.ModelCallContextID}) + return modelOutputAuthorityFromGetSQLC(row), err } func IntegrationCreateContentBlockTx( @@ -435,38 +431,103 @@ func IntegrationCreateContentBlockTx( db dbsqlc.DBTX, input CreateContentBlockInput, ) (ContentBlockRecord, error) { - return createContentBlockTx(ctx, db, input) + metadata, err := input.Metadata.JSON() + if err != nil { + return ContentBlockRecord{}, err + } + var result ContentBlockRecord + err = db.QueryRow(ctx, + `INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,owner_model_output_id,owner_tool_call_result_id,ordinal,block_kind,text_content,structured_data,artifact_id,tool_call_id,exclude_from_model_context,metadata,created_at) + VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,agent_payload_created_at($14,$1,$2, + coalesce($3::uuid,$4::uuid,$5::uuid),false)) RETURNING id,created_at`, + input.AgentID, + input.OwnerKind, + storeutil.IDFromNil(input.OwnerAgentInputID), + storeutil.IDFromNil(input.OwnerModelOutputID), + storeutil.IDFromNil(input.OwnerToolCallResultID), + input.Ordinal, + input.BlockKind, + storeutil.TextFromEmpty(input.TextContent), + sqlcRawMessageFromEmpty(input.StructuredData), + storeutil.IDFromNil(input.ArtifactID), + storeutil.IDFromNil(input.ToolCallID), + input.ExcludeFromModelContext, + metadata, + input.ProjectID). + Scan(&result.ID, &result.CreatedAt) + return result, err } func IntegrationAppendTypedAgentEventTx( ctx context.Context, - txNotifications *notifications.TxNotifications, + _ *notifications.TxNotifications, tx pgx.Tx, input AppendTypedAgentEventInput, ) (TypedAgentEventRecord, error) { - return appendTypedAgentEventTx(ctx, txNotifications, tx, input) + id := input.ID + if id == uuid.Nil { + id = uuid.New() + } + result := TypedAgentEventRecord{ + TurnID: input.TurnID, + IsOpeningEvent: input.IsOpeningEvent, + AgentInputID: input.AgentInputID, + ModelOutputID: input.ModelOutputID, + ToolCallResultID: input.ToolCallResultID, + ContextCheckpointID: input.ContextCheckpointID, + Event: events.Event{ + ID: id, + AgentID: input.AgentID, + Kind: input.Kind, + IdempotencyKey: input.IdempotencyKey, + }, + } + kind := string(input.Kind) + if input.ToolCallResultID != uuid.Nil { + kind = "tool_call_result" + } + err := tx.QueryRow(ctx, + `WITH allocated AS (UPDATE agents SET next_event_sequence=next_event_sequence+1 WHERE id=$1 AND project_id=$2 RETURNING next_event_sequence-1 AS sequence) + INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,agent_input_id, + model_output_id,context_checkpoint_id,tool_call_result_id,is_opening_event,idempotency_key,created_at) + SELECT $3,$1,$4,sequence,$5,$6,$7,$8,$9,$10,$11,agent_payload_created_at($2,$1,$12, + coalesce($6::uuid,$7::uuid,$8::uuid,$9::uuid),true) FROM allocated RETURNING sequence,created_at`, + input.AgentID, + input.ProjectID, + id, + input.TurnID, + input.Kind, + storeutil.IDFromNil(input.AgentInputID), + storeutil.IDFromNil(input.ModelOutputID), + storeutil.IDFromNil(input.ContextCheckpointID), + storeutil.IDFromNil(input.ToolCallResultID), + input.IsOpeningEvent, + storeutil.TextFromEmpty(input.IdempotencyKey), + kind). + Scan(&result.Event.Sequence, &result.Event.At) + return result, err } func IntegrationUpdateAgentTurnLatestEventQuery( ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, turnID, latestEventID, latestSemanticEventID uuid.UUID, + db dbsqlc.DBTX, + projectID, agentID, turnID, latest, semantic uuid.UUID, ) error { - return updateAgentTurnLatestEventQuery( + _, err := db.Exec( ctx, - qtx, + `UPDATE agent_turns t SET latest_event_id=$4,latest_semantic_event_id=coalesce($5,latest_semantic_event_id) FROM agents a WHERE a.id=t.agent_id AND a.project_id=$1 AND t.agent_id=$2 AND t.id=$3`, projectID, agentID, turnID, - latestEventID, - latestSemanticEventID, + latest, + storeutil.IDFromNil(semantic), ) + return err } func IntegrationActivateAgentConfigTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input ActivateAgentConfigInput, ) (AgentConfigChangeRecord, error) { @@ -477,13 +538,13 @@ func IntegrationActivateAgentConfigTx( if err := lockAgentConfigForUseTx(ctx, qtx, config); err != nil { return AgentConfigChangeRecord{}, err } - if err := lockAgentForConfigActivationTx(ctx, qtx, input); err != nil { + if err := lockAgentForConfigActivationTx(ctx, unit, input); err != nil { return AgentConfigChangeRecord{}, err } if err := authorizeAgentConfigChangeTx(ctx, qtx, input); err != nil { return AgentConfigChangeRecord{}, err } - return activateLockedAuthorizedAgentConfigTx(ctx, txNotifications, tx, qtx, input) + return activateLockedAuthorizedAgentConfigTx(ctx, unit, qtx, input) } func (s *Store) IntegrationCompleteDaemonProcessAction( @@ -525,16 +586,6 @@ func IntegrationMachineStillUnreachableForToolExpiryTx( return machineStillUnreachableForToolExpiryTx(ctx, qtx, orgID, machineID, fallbackAt, graceSeconds) } -func IntegrationAgentRuntimeLockRecordFromCancelSQLC( - row dbsqlc.AgentRuntimeLock, -) AgentRuntimeLockRecord { - return agentRuntimeLockRecordFromSQLC(row) -} - -func IntegrationToolCallRecordFromInsertSQLC(row dbsqlc.InsertToolCallRow) ToolCallRecord { - return toolCallRecordFromInsertSQLC(row) -} - func IntegrationAgentMachineBindingRecordFromSQLC( row dbsqlc.GetAgentMachineBindingByMachineRow, ) AgentMachineBindingRecord { @@ -560,11 +611,59 @@ func IntegrationGetToolCallTx( func IntegrationAppendToolResultEventTx( ctx context.Context, - txNotifications *notifications.TxNotifications, + notifications *notifications.TxNotifications, tx pgx.Tx, record ToolCallRecord, ) (events.Event, error) { - return appendToolResultEventTx(ctx, txNotifications, tx, record, nil) + var resultID uuid.UUID + err := tx.QueryRow(ctx, + `INSERT INTO tool_call_results(agent_id,tool_call_id,outcome,completed_at) VALUES($1,$2,$3,statement_timestamp()) RETURNING id`, + record.AgentID, + record.ID, + record.Outcome). + Scan(&resultID) + if err != nil { + return events.Event{}, err + } + blocks, err := parseToolResultContentBlocks(record.ResultContentParts) + if err != nil { + return events.Event{}, err + } + for _, block := range blocks { + block.ProjectID = record.ProjectID + block.AgentID = record.AgentID + block.OwnerKind = ContentBlockOwnerToolCallResult + block.OwnerToolCallResultID = resultID + if _, err = IntegrationCreateContentBlockTx(ctx, tx, block); err != nil { + return events.Event{}, err + } + } + event, err := IntegrationAppendTypedAgentEventTx( + ctx, + notifications, + tx, + AppendTypedAgentEventInput{ + ProjectID: record.ProjectID, + AgentID: record.AgentID, + TurnID: record.TurnID, + Kind: events.KindToolResult, + ToolCallResultID: resultID, + IdempotencyKey: "tool_result:" + record.ID.String(), + }, + ) + if err != nil { + return events.Event{}, err + } + err = IntegrationUpdateAgentTurnLatestEventQuery( + ctx, + tx, + record.ProjectID, + record.AgentID, + record.TurnID, + event.Event.ID, + event.Event.ID, + ) + return event.Event, err } func (s *Store) ArchiveIdleAgentsAsOf( @@ -585,3 +684,84 @@ func (s *Store) ArchiveIdleAgentCandidateAsOf( ) return archived, err } + +func IntegrationAdmitInputs( + ctx context.Context, + unit *agentexecution.Unit, + projectID, agentID uuid.UUID, +) (AdmittedAgentInputTurn, error) { + h, err := unit.Handle(projectID, agentID) + if err != nil { + return AdmittedAgentInputTurn{}, err + } + admitted, err := h.AdmitInputs(ctx) + if err != nil { + return AdmittedAgentInputTurn{}, err + } + if err := applyAdmissionDestination(ctx, unit, projectID, agentID, admitted); err != nil { + return AdmittedAgentInputTurn{}, err + } + return shapeAdmission(ctx, dbsqlc.New(unit.DB()), projectID, agentID, admitted) +} + +type AppendTypedAgentEventInput struct { + ID uuid.UUID + ProjectID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + IsOpeningEvent bool + Kind events.Kind + IdempotencyKey string + AgentInputID uuid.UUID + ModelOutputID uuid.UUID + ToolCallResultID uuid.UUID + ContextCheckpointID uuid.UUID +} + +type ContentBlockRecord struct { + ID uuid.UUID + ProjectID uuid.UUID + AgentID uuid.UUID + OwnerKind ContentBlockOwnerKind + OwnerAgentInputID uuid.UUID + OwnerModelOutputID uuid.UUID + OwnerToolCallResultID uuid.UUID + Ordinal int32 + BlockKind ContentBlockKind + TextContent string + StructuredData json.RawMessage + ArtifactID uuid.UUID + ToolCallID uuid.UUID + CreatedAt time.Time +} + +func sqlcRawMessageFromEmpty(value json.RawMessage) *json.RawMessage { + if len(value) == 0 { + return nil + } + return &value +} + +func toolCallResultAuthorityFromGetSQLC( + row dbsqlc.GetToolCallResultByToolCallRow, +) ToolCallResultAuthorityRecord { + return ToolCallResultAuthorityRecord{ + ID: row.ID, + ProjectID: row.ProjectID, + AgentID: row.AgentID, + TurnID: row.TurnID, + ToolCallID: row.ToolCallID, + Outcome: ToolResultOutcome(row.Outcome), + CompletedAt: row.CompletedAt, + } +} + +type ToolCallResultAuthorityRecord struct { + ID uuid.UUID + ProjectID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + ToolCallID uuid.UUID + Outcome ToolResultOutcome + CompletedAt time.Time +} diff --git a/internal/storage/executionstore/inbox_input_integration_test.go b/internal/storage/executionstore/inbox_input_integration_test.go index 07946b3c1..f2c44620a 100644 --- a/internal/storage/executionstore/inbox_input_integration_test.go +++ b/internal/storage/executionstore/inbox_input_integration_test.go @@ -42,7 +42,10 @@ func inboxInputIntegration( integration, err := f.store.Integrations().CreateIntegration( f.ctx, integrationstore.SaveIntegrationInput{ - OrgID: testOrgID, ProjectID: testProjectID, Name: "github", IntegrationKind: integrationdefinition.GitHubPR, + OrgID: testOrgID, + ProjectID: testProjectID, + Name: "github", + IntegrationKind: integrationdefinition.GitHubPR, }, ) require.NoError(t, err) @@ -54,16 +57,13 @@ func inboxInputIntegration( ProjectID: testProjectID, IntegrationID: integration.ID, InstalledByUserID: f.user.ID, - Provider: provider, ProviderTenantID: "123", ProviderAccountRef: "456", CredentialAppID: 123, - CredentialSecretID: secret.ID, CredentialVersionID: secret.CurrentVersionID, ExpectedSetupRevision: integration.SetupRevision, - ProviderIdentity: json.RawMessage(`{"bot_user_id":789,"bot_login":"kernel-test[bot]"}`), }, ) @@ -87,18 +87,21 @@ func inboxInputPlan( integration integrationstore.IntegrationRecord, event string, ) executionstore.InboxInputRecipient { - return executionstore.InboxInputRecipient{AgentID: agentID, Input: executionstore.CreateAgentContentInputInput{ - Origin: &executionstore.AgentInputOrigin{ - IntegrationID: integration.ID, - Address: integrationstore.ConversationAddress{Kind: "pull_request", Ref: "123#42"}, + return executionstore.InboxInputRecipient{ + AgentID: agentID, + Input: executionstore.CreateAgentContentInputInput{ + Origin: &executionstore.AgentInputOrigin{ + IntegrationID: integration.ID, + Address: integrationstore.ConversationAddress{Kind: "pull_request", Ref: "123#42"}, + }, + Actor: mustIntegrationActorParams(t, integration, "participant"), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"Please address this review"}]`), + Metadata: json.RawMessage(`{"source":"review"}`), + IdempotencyKey: event, + DeliveryMode: executionstore.DeliveryModeSteering, + CancelOpenInteractions: true, }, - Actor: mustIntegrationActorParams(t, integration, "participant"), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"Please address this review"}]`), - Metadata: json.RawMessage(`{"source":"review"}`), - IdempotencyKey: event, - DeliveryMode: executionstore.DeliveryModeSteering, - CancelOpenInteractions: true, - }} + } } func freezeInboxInput( @@ -131,7 +134,8 @@ func freezeInboxInput( ) require.NoError(t, err) require.True(t, found) - integration, err := f.store.Integrations().GetIntegration(f.ctx, testProjectID, recipient.Input.Origin.IntegrationID) + integration, err := f.store.Integrations(). + GetIntegration(f.ctx, testProjectID, recipient.Input.Origin.IntegrationID) require.NoError(t, err) switch integration.Provider { case integrationdefinition.ProviderSlack: @@ -163,15 +167,20 @@ func inboxInputArtifacts(recipient executionstore.InboxInputRecipient) []artifac return artifacts } -func withInboxFile(t *testing.T, recipient executionstore.InboxInputRecipient) executionstore.InboxInputRecipient { +func withInboxFile( + t *testing.T, + recipient executionstore.InboxInputRecipient, +) executionstore.InboxInputRecipient { t.Helper() id, err := uuid.NewV7() require.NoError(t, err) recipient.ArtifactIDs = []uuid.UUID{id} - recipient.Files = []executionstore.InboxPlannedFile{{ArtifactID: id, Expected: &artifactstore.PreparedArtifact{ - ID: id, ContentType: "text/plain", Filename: "review.txt", - Digest: blobstore.ContentDigest([]byte("review")), SizeBytes: 6, - }}} + recipient.Files = []executionstore.InboxPlannedFile{ + {ArtifactID: id, Expected: &artifactstore.PreparedArtifact{ + ID: id, ContentType: "text/plain", Filename: "review.txt", + Digest: blobstore.ContentDigest([]byte("review")), SizeBytes: 6, + }}, + } recipient.Input.ContentBlocks = json.RawMessage( `[{"type":"text","text":"Review attached"},{"type":"media_ref","artifact_id":"` + id.String() + `"}]`, ) @@ -209,7 +218,11 @@ func TestInboxInputGitHubCommentsSteerAndCancelAcrossProviders(t *testing.T) { ) require.NoError(t, err) require.True(t, result.Created) - require.ElementsMatch(t, []uuid.UUID{firstPrompt.ID, secondPrompt.ID}, result.CanceledInteractionIDs) + require.ElementsMatch( + t, + []uuid.UUID{firstPrompt.ID, secondPrompt.ID}, + result.CanceledInteractionIDs, + ) for _, id := range result.CanceledInteractionIDs { prompt := f.read(t, id) require.Equal(t, executionstore.AgentInteractionStateCanceled, prompt.State) @@ -226,11 +239,16 @@ func TestInboxInputGitHubCommentsSteerAndCancelAcrossProviders(t *testing.T) { require.NoError(t, err) require.Equal(t, before.CurrentConfigID, after.CurrentConfigID) require.Empty(t, f.activation().subscriptions(t, f.process.AgentID)) - selection, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) + selection, err := f.store.Execution(). + GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) require.Equal( t, - executionstore.InteractionSelection{AutoSelect: true, HandlerKey: "other", IntegrationTargetID: f.b.ID}, + executionstore.InteractionSelection{ + AutoSelect: true, + HandlerKey: "other", + IntegrationTargetID: f.b.ID, + }, selection, "receipt does not change the current turn destination before input admission", ) @@ -238,13 +256,15 @@ func TestInboxInputGitHubCommentsSteerAndCancelAcrossProviders(t *testing.T) { selected := f.selectOrigin(t, f.a.ID) newPrompt := createQuestionInteractionForTest(t, f.ctx, f.process, calls[2]) duplicate := freezeInboxInput(t, f.activation(), recipient, "other-callback", time.Minute) - replayed, err := f.store.Execution().AdmitInboxInputRecipient(f.ctx, duplicate.Lease(), "recipient", nil) + replayed, err := f.store.Execution(). + AdmitInboxInputRecipient(f.ctx, duplicate.Lease(), "recipient", nil) require.NoError(t, err) require.False(t, replayed.Created) require.Equal(t, result.AgentInput.ID, replayed.AgentInput.ID) require.Empty(t, replayed.CanceledInteractionIDs) require.Equal(t, executionstore.AgentInteractionStateOpen, f.read(t, newPrompt.ID).State) - selection, err = f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) + selection, err = f.store.Execution(). + GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) require.Equal(t, selected, selection) require.NoError( @@ -256,7 +276,8 @@ func TestInboxInputGitHubCommentsSteerAndCancelAcrossProviders(t *testing.T) { recipient.Input.DeliveryMode = executionstore.DeliveryModeQueued recipient.Input.CancelOpenInteractions = false queued := freezeInboxInput(t, f.activation(), recipient, "commit-callback", time.Minute) - result, err = f.store.Execution().AdmitInboxInputRecipient(f.ctx, queued.Lease(), "recipient", nil) + result, err = f.store.Execution(). + AdmitInboxInputRecipient(f.ctx, queued.Lease(), "recipient", nil) require.NoError(t, err) require.Equal(t, executionstore.DeliveryModeQueued, result.AgentInput.DeliveryMode) require.Empty(t, result.CanceledInteractionIDs) @@ -264,7 +285,8 @@ func TestInboxInputGitHubCommentsSteerAndCancelAcrossProviders(t *testing.T) { recipient.Input.DeliveryMode = executionstore.DeliveryModeSteering recipient.Input.IdempotencyKey = "steer-without-cancel" steering := freezeInboxInput(t, f.activation(), recipient, "steering-callback", time.Minute) - result, err = f.store.Execution().AdmitInboxInputRecipient(f.ctx, steering.Lease(), "recipient", nil) + result, err = f.store.Execution(). + AdmitInboxInputRecipient(f.ctx, steering.Lease(), "recipient", nil) require.NoError(t, err) require.Empty(t, result.CanceledInteractionIDs) require.Equal(t, executionstore.AgentInteractionStateOpen, f.read(t, newPrompt.ID).State) @@ -275,7 +297,8 @@ func TestInboxInputGitHubCommentsSteerAndCancelAcrossProviders(t *testing.T) { func TestInboxInputConcurrentMediaAndCompletedReplay(t *testing.T) { t.Parallel() f := newIntegrationActivationFixture(t) - agent, err := f.store.Execution().LaunchAgent(f.ctx, f.launchInput(f.profile.CurrentConfigID, "existing-recipient")) + agent, err := f.store.Execution(). + LaunchAgent(f.ctx, f.launchInput(f.profile.CurrentConfigID, "existing-recipient")) require.NoError(t, err) recipient := withInboxFile(t, inboxInputPlan(t, agent.Agent.ID, f.integration, "message:file")) recipient.Input.Origin.Address = integrationstore.ConversationAddress{Kind: "thread", Ref: "C123:111.222"} @@ -283,7 +306,8 @@ func TestInboxInputConcurrentMediaAndCompletedReplay(t *testing.T) { second := freezeInboxInput(t, f, recipient, "file-two", time.Minute) start := make(chan struct{}) - admit := func(receipt integrationstore.IntegrationInboxRecord) func() (executionstore.InboxInputResult, error) { + admit := func(receipt integrationstore.IntegrationInboxRecord) func() (executionstore.InboxInputResult, + error) { return func() (executionstore.InboxInputResult, error) { <-start return f.store.Execution().AdmitInboxInputRecipient( @@ -293,7 +317,15 @@ func TestInboxInputConcurrentMediaAndCompletedReplay(t *testing.T) { } one, two := integrationdb.RunAsync(admit(first)), integrationdb.RunAsync(admit(second)) close(start) - a, b := integrationdb.AwaitSuccess(t, one, "first event"), integrationdb.AwaitSuccess(t, two, "duplicate event") + a, b := integrationdb.AwaitSuccess( + t, + one, + "first event", + ), integrationdb.AwaitSuccess( + t, + two, + "duplicate event", + ) require.NotEqual(t, a.Created, b.Created) require.Equal(t, a.AgentInput.ID, b.AgentInput.ID) for _, table := range []string{"artifacts", "integration_targets"} { @@ -344,7 +376,11 @@ func TestInboxInputExpiredLeaseRollsBackOriginMediaAndCancellation(t *testing.T) 10*time.Millisecond, ) require.NoError(t, blocker.Commit(f.ctx)) - require.ErrorIs(t, integrationdb.Await(t, done, "expired input").Err, integrationstore.ErrIntegrationInboxLeaseLost) + require.ErrorIs( + t, + integrationdb.Await(t, done, "expired input").Err, + integrationstore.ErrIntegrationInboxLeaseLost, + ) require.Equal(t, executionstore.AgentInteractionStateOpen, f.read(t, prompt.ID).State) current, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) @@ -352,7 +388,8 @@ func TestInboxInputExpiredLeaseRollsBackOriginMediaAndCancellation(t *testing.T) var count int require.NoError( t, - f.store.pool.QueryRow(f.ctx, `SELECT count(*) FROM artifacts WHERE id=$1`, recipient.ArtifactIDs[0]).Scan(&count), + f.store.pool.QueryRow(f.ctx, `SELECT count(*) FROM artifacts WHERE id=$1`, recipient.ArtifactIDs[0]). + Scan(&count), ) require.Zero(t, count) require.NoError( @@ -391,7 +428,8 @@ func TestInboxInputExpiredLeaseRollsBackOriginMediaAndCancellation(t *testing.T) func TestInboxInputIntegrationGateBeforeReceiptAndAgent(t *testing.T) { t.Parallel() f := newIntegrationActivationFixture(t) - agent, err := f.store.Execution().LaunchAgent(f.ctx, f.launchInput(f.profile.CurrentConfigID, "lock-recipient")) + agent, err := f.store.Execution(). + LaunchAgent(f.ctx, f.launchInput(f.profile.CurrentConfigID, "lock-recipient")) require.NoError(t, err) recipient := inboxInputPlan(t, agent.Agent.ID, f.integration, "message:lock") recipient.Input.Origin.Address = integrationstore.ConversationAddress{Kind: "thread", Ref: "C123:111.222"} @@ -418,7 +456,10 @@ func TestInboxInputIntegrationGateBeforeReceiptAndAgent(t *testing.T) { dbsqlc.LockIntegrationInboxReceiptParams{ProjectID: testProjectID, ID: receipt.ID}, ) require.NoError(t, err) - _, err = q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: agent.Agent.ID}) + _, err = q.LockAgentInProject( + ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: agent.Agent.ID}, + ) require.NoError(t, err) require.NoError(t, blocker.Commit(f.ctx)) require.True(t, integrationdb.AwaitSuccess(t, done, "ordered input admission").Created) @@ -434,26 +475,57 @@ func TestInboxInputSelectionChangesOnlyOnContentAdmission(t *testing.T) { f := newIntegrationInteractionFixture(t) before := f.selectOrigin(t, f.a.ID) recipient := inboxInputPlan(t, f.process.AgentID, f.otherIntegration, "handler-origin") - recipient.Input.Origin.Address = integrationstore.ConversationAddress{Kind: "thread", Ref: f.b.ScopeRef} + recipient.Input.Origin.Address = integrationstore.ConversationAddress{ + Kind: "thread", + Ref: f.b.ScopeRef, + } recipient.Input.DeliveryMode = mode recipient.Input.CancelOpenInteractions = false receipt := freezeInboxInput(t, f.activation(), recipient, "handler-origin", time.Minute) - result, err := f.store.Execution().AdmitInboxInputRecipient(f.ctx, receipt.Lease(), "recipient", nil) + result, err := f.store.Execution(). + AdmitInboxInputRecipient(f.ctx, receipt.Lease(), "recipient", nil) require.NoError(t, err) require.True(t, result.Created) - selected, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) + selected, err := f.store.Execution(). + GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) - require.Equal(t, before, selected, "receiving an input must not change the current turn's destination") + require.Equal( + t, + before, + selected, + "receiving an input must not change the current turn's destination", + ) admitted, found := admitNextAgentInputAndOpenTurnForTest( t, f.ctx, f.store, testProjectID, f.process.AgentID, f.process.Lock.ID) require.True(t, found) require.Len(t, admitted.Inputs, 1) require.Equal(t, result.AgentInput.ID, admitted.Inputs[0].ID) - selected, err = f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) + selected, err = f.store.Execution(). + GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) require.Equal(t, "other", selected.HandlerKey) require.Equal(t, f.b.ID, selected.IntegrationTargetID) + prepared, err := f.store.Execution(). + PrepareNormalModelCall(f.ctx, + executionstore.PrepareNormalModelCallInput{ProjectID: testProjectID, + AgentID: f.process.AgentID, + RuntimeLockID: f.process.Lock.ID, + OpeningInputIDs: []uuid.UUID{admitted.Inputs[0].ID}}) + require.NoError(t, err) + contextID := prepared.Claim.Context.ID + createModelOutputEventForTurnTest( + t, + f.ctx, + f.process, + admitted.Turn.ID, + contextID, + "origin-done", + "end_turn", + "", + f.process.Now, + ) + dashboard, _, _, err := f.store.Execution().CreateAgentContentInput( f.ctx, executionstore.CreateAgentContentInputInput{ ProjectID: testProjectID, AgentID: f.process.AgentID, Actor: mustOmnaraActorParams(t, f.user.ID), @@ -474,7 +546,8 @@ func TestInboxInputSelectionChangesOnlyOnContentAdmission(t *testing.T) { require.Equal(t, executionstore.InteractionSelection{AutoSelect: true}, after) current := f.selectOrigin(t, f.a.ID) - replayed, err := f.store.Execution().AdmitInboxInputRecipient(f.ctx, receipt.Lease(), "recipient", nil) + replayed, err := f.store.Execution(). + AdmitInboxInputRecipient(f.ctx, receipt.Lease(), "recipient", nil) require.NoError(t, err) require.False(t, replayed.Created) _, found = admitNextAgentInputAndOpenTurnForTest( @@ -508,7 +581,12 @@ func TestInboxSubscriptionInputDoesNotGrantInteractionDestination(t *testing.T) require.True(t, result.Created) selected, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) - require.Equal(t, before, selected, "subscription receipt must wait for content admission to change selection") + require.Equal( + t, + before, + selected, + "subscription receipt must wait for content admission to change selection", + ) admitted, found := admitNextAgentInputAndOpenTurnForTest( t, f.ctx, f.store, testProjectID, f.process.AgentID, f.process.Lock.ID) require.True(t, found) @@ -593,7 +671,10 @@ func TestInboxInputInvalidArtifactsOrActorLeavesNoAdmission(t *testing.T) { LaunchAgent(f.ctx, f.launchInput(f.profile.CurrentConfigID, "invalid-recipient")) require.NoError(t, err) recipient := withInboxFile(t, inboxInputPlan(t, agent.Agent.ID, f.integration, "invalid-file")) - recipient.Input.Origin.Address = integrationstore.ConversationAddress{Kind: "thread", Ref: "C123:111.222"} + recipient.Input.Origin.Address = integrationstore.ConversationAddress{ + Kind: "thread", + Ref: "C123:111.222", + } switch scenario { case "external-actor": recipient.Input.Actor.Provider = executionstore.ActorProviderExternal @@ -610,7 +691,8 @@ func TestInboxInputInvalidArtifactsOrActorLeavesNoAdmission(t *testing.T) { if scenario == "substituted-artifact" { artifacts[0].ID = uuid.Must(uuid.NewV7()) } - _, err = f.store.Execution().AdmitInboxInputRecipient(f.ctx, receipt.Lease(), "recipient", artifacts) + _, err = f.store.Execution(). + AdmitInboxInputRecipient(f.ctx, receipt.Lease(), "recipient", artifacts) require.Error(t, err) for _, check := range []struct { query string @@ -642,7 +724,10 @@ func TestInboxMessageSiblingsConcurrentAndDelayedFiles(t *testing.T) { t.Parallel() f := newIntegrationInteractionFixture(t) base := inboxInputPlan(t, f.process.AgentID, f.integration, "text-callback") - base.Input.Origin.Address = integrationstore.ConversationAddress{Kind: "thread", Ref: "C123:111.222"} + base.Input.Origin.Address = integrationstore.ConversationAddress{ + Kind: "thread", + Ref: "C123:111.222", + } base.Sibling = &executionstore.InboxMessageSibling{Key: "file-callback"} files := withInboxFile(t, base) files.Input.IdempotencyKey = "file-callback" diff --git a/internal/storage/executionstore/inbox_input_store.go b/internal/storage/executionstore/inbox_input_store.go index bbdb8fa41..1c2ad57be 100644 --- a/internal/storage/executionstore/inbox_input_store.go +++ b/internal/storage/executionstore/inbox_input_store.go @@ -6,10 +6,10 @@ import ( "errors" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/integrationdefinition" "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -79,14 +79,15 @@ func (s *Store) admitInboxInputRecipientOnce( } return replayInboxInput(ctx, s.q, recipient, outcome) } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return InboxInputResult{}, err } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() work, err := s.integrations.LockIntegrationInboxLeaseTx(ctx, tx, lease) if err != nil { - _ = tx.Rollback(ctx) + _ = unit.Rollback(ctx) if outcome, readErr := resolveInboxInputOutcome( ctx, s.q, recipient, ); readErr == nil && outcome.Outcome != InboxRecipientPending { @@ -108,9 +109,9 @@ func (s *Store) admitInboxInputRecipientOnce( ); err != nil { return InboxInputResult{}, err } - if err := lifecyclelock.Agents(ctx, tx, []lifecyclelock.AgentRef{{ + if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ ProjectID: lease.ProjectID, AgentID: recipient.AgentID, - }}); err != nil { + }}, agentexecution.IngressAuthority{}); err != nil { return InboxInputResult{}, err } outcome, err := resolveInboxInputOutcome(ctx, q, recipient) @@ -147,11 +148,10 @@ func (s *Store) admitInboxInputRecipientOnce( return InboxInputResult{}, err } } - notifications := s.newTxNotifications() + result, err := s.admitOriginContentTx( ctx, - tx, - notifications, + unit, recipient.Input, blocks, artifacts, @@ -162,7 +162,7 @@ func (s *Store) admitInboxInputRecipientOnce( if err := work.CheckLease(ctx); err != nil { return InboxInputResult{}, err } - if err := s.commitTxWithNotifications(ctx, tx, notifications, "admit inbox input recipient"); err != nil { + if err := unit.Commit(ctx, "admit inbox input recipient"); err != nil { return InboxInputResult{}, err } return result, nil diff --git a/internal/storage/executionstore/inbox_subscription.go b/internal/storage/executionstore/inbox_subscription.go index af04100a8..bd0d783f1 100644 --- a/internal/storage/executionstore/inbox_subscription.go +++ b/internal/storage/executionstore/inbox_subscription.go @@ -3,12 +3,11 @@ package executionstore import ( "context" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) -func validateInboxSubscriptionTx(ctx context.Context, tx pgx.Tx, recipient InboxInputRecipient) error { +func validateInboxSubscriptionTx(ctx context.Context, tx dbsqlc.DBTX, recipient InboxInputRecipient) error { q := dbsqlc.New(tx) for _, address := range recipient.Subscription.Alternatives { allowed, err := q.HasIntegrationSubscription(ctx, dbsqlc.HasIntegrationSubscriptionParams{ diff --git a/internal/storage/executionstore/integration_config_activation.go b/internal/storage/executionstore/integration_config_activation.go index 1bea54327..b49648e39 100644 --- a/internal/storage/executionstore/integration_config_activation.go +++ b/internal/storage/executionstore/integration_config_activation.go @@ -64,7 +64,7 @@ func configChangeReplayExistsTx(ctx context.Context, q *dbsqlc.Queries, input Ch func lockConfigChangeIntegrationsTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, q *dbsqlc.Queries, input ChangeAgentConfigInput, ) error { diff --git a/internal/storage/executionstore/integration_handler_targets.go b/internal/storage/executionstore/integration_handler_targets.go index 7726d06bb..4b32264a1 100644 --- a/internal/storage/executionstore/integration_handler_targets.go +++ b/internal/storage/executionstore/integration_handler_targets.go @@ -6,8 +6,8 @@ import ( "strings" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) type integrationConversation struct { @@ -17,7 +17,7 @@ type integrationConversation struct { func lockIntegrationConversationsTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID uuid.UUID, origins ...AgentInputOrigin, ) error { diff --git a/internal/storage/executionstore/integration_input_origin.go b/internal/storage/executionstore/integration_input_origin.go index 4f6254d55..5b69c7781 100644 --- a/internal/storage/executionstore/integration_input_origin.go +++ b/internal/storage/executionstore/integration_input_origin.go @@ -7,10 +7,9 @@ import ( "strings" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -60,7 +59,7 @@ func prepareOriginContentInput( func (s *Store) resolveInputOriginTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, input CreateAgentContentInputInput, ) (CreateAgentContentInputInput, integrationstore.IntegrationRecord, error) { integration, err := s.integrations.GetIntegrationByIDTx(ctx, tx, input.Origin.IntegrationID) @@ -90,16 +89,16 @@ func validateIntegrationInputActor(integration integrationstore.IntegrationRecor func (s *Store) admitOriginContentTx( ctx context.Context, - tx pgx.Tx, - notifications *notifications.TxNotifications, + unit *agentexecution.Unit, input CreateAgentContentInputInput, blocks []CreateContentBlockInput, artifacts []artifactstore.PreparedArtifact, ) (InboxInputResult, error) { - if err := lifecyclelock.Agents( + tx := unit.DB() + if err := unit.LockAgentRefs( ctx, - tx, []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, AgentID: input.AgentID}}, + agentexecution.IngressAuthority{}, ); err != nil { return InboxInputResult{}, err } @@ -125,7 +124,7 @@ func (s *Store) admitOriginContentTx( if err != nil { return InboxInputResult{}, err } - created, err := createAgentContentInputTx(ctx, notifications, tx, dbsqlc.New(tx), input, blocks) + created, err := createAgentContentInputTx(ctx, unit, dbsqlc.New(tx), input, blocks) if err != nil { return InboxInputResult{}, err } @@ -141,7 +140,7 @@ func (s *Store) admitOriginContentTx( func (s *Store) originContentReplayTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, input CreateAgentContentInputInput, ) (InboxInputResult, bool, error) { q := dbsqlc.New(tx) diff --git a/internal/storage/executionstore/integration_interactions_integration_test.go b/internal/storage/executionstore/integration_interactions_integration_test.go index bb905f603..7fd12d788 100644 --- a/internal/storage/executionstore/integration_interactions_integration_test.go +++ b/internal/storage/executionstore/integration_interactions_integration_test.go @@ -18,6 +18,7 @@ import ( "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -81,7 +82,11 @@ func newIntegrationInteractionFixture(t *testing.T) integrationInteractionFixtur ) require.NoError( t, - lifecyclelock.Agents(ctx, tx, []lifecyclelock.AgentRef{{ProjectID: testProjectID, AgentID: f.process.AgentID}}), + lifecyclelock.Agents( + ctx, + tx, + []lifecyclelock.AgentRef{{ProjectID: testProjectID, AgentID: f.process.AgentID}}, + ), ) for _, target := range []integrationstore.IntegrationTargetRecord{f.a, f.b} { require.NoError( @@ -225,16 +230,18 @@ func (f integrationInteractionFixture) selectOrigin( targetID uuid.UUID, ) executionstore.InteractionSelection { t.Helper() - tx := integrationdb.BeginTx(t, f.ctx, f.store.pool) - _, err := dbsqlc.New(tx).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ + txUnit, unitErr := f.store.Execution().IntegrationBeginUnit(f.ctx) + require.NoError(t, unitErr) + defer func() { _ = txUnit.Rollback(f.ctx) }() + _, err := txUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, - }) + }, agentexecution.ExternalAuthority{}) require.NoError(t, err) selection, err := executionstore.IntegrationSelectInteractionDestinationForOriginTx( - f.ctx, tx, testProjectID, f.process.AgentID, targetID, + f.ctx, txUnit, testProjectID, f.process.AgentID, targetID, ) require.NoError(t, err) - require.NoError(t, tx.Commit(f.ctx)) + require.NoError(t, txUnit.Commit(f.ctx, "fixture admission")) stored, err := f.store.Execution(). GetInteractionSelection(f.ctx, testProjectID, f.process.AgentID) require.NoError(t, err) @@ -434,7 +441,11 @@ func TestIntegrationInteractionsRevocationPreservesDashboardAndSnapshot(t *testi case "integration disconnected": f.disable(t) case "integration deleted": - _, err := f.store.pool.Exec(f.ctx, `UPDATE integrations SET deleted_at=now() WHERE id=$1`, f.integration.ID) + _, err := f.store.pool.Exec( + f.ctx, + `UPDATE integrations SET deleted_at=now() WHERE id=$1`, + f.integration.ID, + ) require.NoError(t, err) case "assignment removed": _, err := f.store.pool.Exec( @@ -457,17 +468,19 @@ func TestIntegrationInteractionsRevocationPreservesDashboardAndSnapshot(t *testi require.NoError(t, err) } - tx := integrationdb.BeginTx(t, f.ctx, f.store.pool) - _, err := dbsqlc.New(tx).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ + unit, err := f.store.Execution().IntegrationBeginUnit(f.ctx) + require.NoError(t, err) + defer func() { _ = unit.Rollback(f.ctx) }() + _, err = dbsqlc.New(unit.DB()).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, }) require.NoError(t, err) - selection, err := f.store.Execution().ReconcileInteractionSelectionTx( - f.ctx, tx, testProjectID, f.process.AgentID, + selection, err := f.store.Execution().ReconcileInteractionSelectionInUnit( + f.ctx, unit, testProjectID, f.process.AgentID, ) require.NoError(t, err) require.Equal(t, executionstore.InteractionSelection{AutoSelect: true}, selection) - require.NoError(t, tx.Commit(f.ctx)) + require.NoError(t, unit.Commit(f.ctx, "selection reconciliation")) _, err = f.store.Execution().GetAgentInteractionForPresentation( f.ctx, testProjectID, f.process.AgentID, question.ID, ) @@ -518,18 +531,20 @@ func TestIntegrationInteractionsRejectForeignCallbackAndTarget(t *testing.T) { } otherAgent := mustCreateAgent(t, f.ctx, f.store) foreign := f.target(t, otherAgent, "C123:111.222") - tx := integrationdb.BeginTx(t, f.ctx, f.store.pool) - _, err := dbsqlc.New(tx).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ + txUnit, unitErr := f.store.Execution().IntegrationBeginUnit(f.ctx) + require.NoError(t, unitErr) + defer func() { _ = txUnit.Rollback(f.ctx) }() + _, err := txUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, - }) + }, agentexecution.ExternalAuthority{}) require.NoError(t, err) selection, err := executionstore.IntegrationSelectInteractionDestinationForOriginTx( - f.ctx, tx, testProjectID, f.process.AgentID, foreign.ID, + f.ctx, txUnit, testProjectID, f.process.AgentID, foreign.ID, ) require.NoError(t, err) require.Equal(t, executionstore.InteractionSelection{AutoSelect: true}, selection, "another agent's target is ineligible and must clear automatic selection") - require.NoError(t, tx.Commit(f.ctx)) + require.NoError(t, txUnit.Commit(f.ctx, "fixture admission")) _, err = f.store.Execution(). GetAgentInteractionForPresentation(f.ctx, uuid.New(), f.process.AgentID, question.ID) require.ErrorIs(t, err, storeerr.ErrNotFound) @@ -592,10 +607,12 @@ func TestIntegrationInteractionsCaptureUsesLockedCurrentSelectionWithoutIntegrat integrationGate := integrationdb.BeginTx(t, f.ctx, f.store.pool) require.NoError(t, dbsqlc.New(integrationGate).LockIntegrationLifecycleExclusive(f.ctx, dbsqlc.LockIntegrationLifecycleExclusiveParams{IntegrationID: f.integration.ID})) - selectionTx := integrationdb.BeginTx(t, f.ctx, f.store.pool) - _, err = dbsqlc.New(selectionTx).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ + selectionTxUnit, unitErr := f.store.Execution().IntegrationBeginUnit(f.ctx) + require.NoError(t, unitErr) + defer func() { _ = selectionTxUnit.Rollback(f.ctx) }() + _, err = selectionTxUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, - }) + }, agentexecution.ExternalAuthority{}) require.NoError(t, err) done := integrationdb.RunAsync(func() (executionstore.AgentInteractionRecord, error) { return f.store.Execution(). @@ -606,10 +623,10 @@ func TestIntegrationInteractionsCaptureUsesLockedCurrentSelectionWithoutIntegrat }) integrationdb.WaitForNamedLockWaiters(t, f.ctx, f.store.pool, "LockAgentInProject", 1) _, err = executionstore.IntegrationSelectInteractionDestinationForOriginTx( - f.ctx, selectionTx, testProjectID, f.process.AgentID, f.b.ID, + f.ctx, selectionTxUnit, testProjectID, f.process.AgentID, f.b.ID, ) require.NoError(t, err) - require.NoError(t, selectionTx.Commit(f.ctx)) + require.NoError(t, selectionTxUnit.Commit(f.ctx, "fixture admission")) interaction := integrationdb.AwaitSuccess( t, done, @@ -854,8 +871,10 @@ func TestIntegrationInteractionsCallbackRechecksConfigAfterAgentLockWait(t *test definition := f.definition(t, "handler revoked while callback waits", nil) config, err := f.store.Execution().CreateAgentConfig(f.ctx, definition) require.NoError(t, err) - activation := integrationdb.BeginTx(t, f.ctx, f.store.pool) - _, err = dbsqlc.New(activation).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ + activation, err := f.store.Execution().IntegrationBeginUnit(f.ctx) + require.NoError(t, err) + defer func() { _ = activation.Rollback(f.ctx) }() + _, err = dbsqlc.New(activation.DB()).LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, }) require.NoError(t, err) @@ -865,7 +884,7 @@ func TestIntegrationInteractionsCallbackRechecksConfigAfterAgentLockWait(t *test ResolveAgentInteractionFromHandler(f.ctx, doneInput) }) integrationdb.WaitForNamedLockWaiters(t, f.ctx, f.store.pool, "LockAgentInProject", 1) - _, err = activation.Exec( + _, err = activation.DB().Exec( f.ctx, `UPDATE agents SET current_config_id=$3 WHERE project_id=$1 AND id=$2`, testProjectID, @@ -874,9 +893,9 @@ func TestIntegrationInteractionsCallbackRechecksConfigAfterAgentLockWait(t *test ) require.NoError(t, err) _, err = f.store.Execution(). - ReconcileInteractionSelectionTx(f.ctx, activation, testProjectID, f.process.AgentID) + ReconcileInteractionSelectionInUnit(f.ctx, activation, testProjectID, f.process.AgentID) require.NoError(t, err) - require.NoError(t, activation.Commit(f.ctx)) + require.NoError(t, activation.Commit(f.ctx, "config activation")) result := integrationdb.Await(t, done, "callback after config activation") require.ErrorIs(t, result.Err, storeerr.ErrUnauthorized) require.Equal(t, executionstore.AgentInteractionStateOpen, f.read(t, question.ID).State) diff --git a/internal/storage/executionstore/integration_interactions_store.go b/internal/storage/executionstore/integration_interactions_store.go index 67c1d63a2..533df2f06 100644 --- a/internal/storage/executionstore/integration_interactions_store.go +++ b/internal/storage/executionstore/integration_interactions_store.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "errors" - "fmt" "reflect" "github.com/google/uuid" @@ -14,6 +13,7 @@ import ( "github.com/omnara-ai/omnara/internal/jsoncanonical" "github.com/omnara-ai/omnara/internal/publicid" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -36,85 +36,28 @@ func (s *Store) GetInteractionSelection( return interactionSelectionFromRow(row), err } -func resolveInteractionSelectionForOriginTx( +func (s *Store) ReconcileInteractionSelectionInUnit( ctx context.Context, - tx pgx.Tx, - projectID, agentID, configID, originTargetID uuid.UUID, + unit *agentexecution.Unit, + projectID, agentID uuid.UUID, ) (InteractionSelection, error) { - q := dbsqlc.New(tx) - selection := InteractionSelection{AutoSelect: true} - if originTargetID == uuid.Nil { - return selection, nil - } - contract, err := interactionContractForConfig(ctx, q, projectID, configID) - if err != nil { - return InteractionSelection{}, err - } - handlers, err := prepareInteractionHandlers(ctx, q, projectID, contract.InteractionHandlers) - if err != nil { - return InteractionSelection{}, err - } - if len(handlers) == 0 { - return selection, nil - } - target, err := q.GetInteractionDestinationTarget( + if err := unit.LockAgentRefs( ctx, - dbsqlc.GetInteractionDestinationTargetParams{ - ProjectID: projectID, - AgentID: agentID, - TargetID: originTargetID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return selection, nil - } - if err != nil { + []lifecyclelock.AgentRef{{ProjectID: projectID, AgentID: agentID}}, + agentexecution.IngressAuthority{}, + ); err != nil { return InteractionSelection{}, err } - matches := 0 - for key, handler := range handlers { - if target.IntegrationState != string(integrationstore.IntegrationStateActive) || - handler.prepared.IntegrationID != target.IntegrationID { - continue - } - destination, err := resolveInteractionHandlerDestination(ctx, tx, projectID, agentID, key, handler) - if err != nil { - return InteractionSelection{}, err - } - if destination == nil || destination.IntegrationTargetID != originTargetID { - continue - } - - matches++ - selection = InteractionSelection{ - AutoSelect: true, - IntegrationTargetID: originTargetID, - HandlerKey: key, - } - } - if matches != 1 { - selection = InteractionSelection{AutoSelect: true} - } - return selection, nil -} - -func (s *Store) ReconcileInteractionSelectionTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID uuid.UUID, -) (InteractionSelection, error) { - selection, destination, err := loadSelectedInteractionDestinationTx(ctx, tx, projectID, agentID) + selection, destination, err := loadSelectedInteractionDestinationTx(ctx, unit.DB(), projectID, agentID) if err != nil || selection.HandlerKey == "" || destination != nil { return selection, err } - cleared := InteractionSelection{AutoSelect: selection.AutoSelect} - return cleared, writeInteractionSelection( - ctx, - dbsqlc.New(tx), - projectID, - agentID, - cleared, - ) + h, err := unit.Handle(projectID, agentID) + if err != nil { + return InteractionSelection{}, err + } + _, err = h.SetInteractionTarget(ctx, uuid.Nil, "", selection.AutoSelect) + return InteractionSelection{AutoSelect: selection.AutoSelect}, err } func (r *ToolCallReader) ListInteractionHandlers( @@ -128,7 +71,7 @@ func (r *ToolCallReader) ListInteractionHandlers( func listInteractionHandlers( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID uuid.UUID, cursor string, limit int, @@ -308,13 +251,15 @@ func SetInteractionHandlerForToolCall( if input.AutoSelect != nil { selected.AutoSelect = *input.AutoSelect } - if err := writeInteractionSelection( - ctx, - t.q, - t.input.ProjectID, - t.input.AgentID, - selected, - ); err != nil { + h, err := t.unit.Handle(t.input.ProjectID, t.input.AgentID) + if err != nil { + return nil, err + } + if _, + err = h.SetInteractionTarget(ctx, + selected.IntegrationTargetID, + selected.HandlerKey, + selected.AutoSelect); err != nil { return nil, err } result, err := completion(selected) @@ -422,10 +367,13 @@ func prepareInteractionHandlers( } handlers := make(map[uuid.UUID]resolvedInteractionHandler, len(rows)) for _, integration := range rows { - if integration.DeletedAt != nil || integration.State != string(integrationstore.IntegrationStateActive) { + if integration.DeletedAt != nil || + integration.State != string(integrationstore.IntegrationStateActive) { continue } - definition, ok := integrationdefinition.Lookup(integrationdefinition.Kind(integration.IntegrationKind)) + definition, ok := integrationdefinition.Lookup( + integrationdefinition.Kind(integration.IntegrationKind), + ) if !ok || definition.InteractionHandler == nil { continue } @@ -449,31 +397,9 @@ func prepareInteractionHandlers( return result, nil } -func writeInteractionSelection( - ctx context.Context, - q *dbsqlc.Queries, - projectID, agentID uuid.UUID, - selection InteractionSelection, -) error { - changed, err := q.SetInteractionSelection(ctx, dbsqlc.SetInteractionSelectionParams{ - AutoSelect: selection.AutoSelect, - ProjectID: projectID, - AgentID: agentID, - TargetID: storeutil.IDFromNil(selection.IntegrationTargetID), - HandlerKey: storeutil.TextFromEmpty(selection.HandlerKey), - }) - if err != nil { - return err - } - if changed != 1 { - return storeerr.ErrStateTransitionConflict - } - return nil -} - func resolveInteractionHandlerDestination( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID uuid.UUID, key string, handler resolvedInteractionHandler, @@ -504,7 +430,7 @@ func resolveInteractionHandlerDestination( } func selectedInteractionDestination( - ctx context.Context, tx pgx.Tx, projectID, agentID uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, projectID, agentID uuid.UUID, selection InteractionSelection, contract agentconfig.RuntimeContract, ) (*InteractionDestination, error) { capability, ok := contract.InteractionHandlers[selection.HandlerKey] @@ -520,7 +446,14 @@ func selectedInteractionDestination( if !ok { return nil, nil //nolint:nilnil // Handler unavailable. } - destination, err := resolveInteractionHandlerDestination(ctx, tx, projectID, agentID, selection.HandlerKey, handler) + destination, err := resolveInteractionHandlerDestination( + ctx, + tx, + projectID, + agentID, + selection.HandlerKey, + handler, + ) if err != nil || destination == nil { return nil, err } @@ -532,7 +465,7 @@ func selectedInteractionDestination( func loadSelectedInteractionDestinationTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID uuid.UUID, ) (InteractionSelection, *InteractionDestination, error) { q := dbsqlc.New(tx) @@ -565,7 +498,7 @@ func loadSelectedInteractionDestinationTx( // presentation and callbacks recheck live authority afterward. func captureInteractionDestinationTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID uuid.UUID, ) (json.RawMessage, error) { _, destination, err := loadSelectedInteractionDestinationTx(ctx, tx, projectID, agentID) @@ -589,9 +522,12 @@ type ResolveAgentInteractionFromHandlerInput struct { func (s *Store) GetInteractionCallbackIntegrationID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { // Cross-project routing hint only; callers must verify the callback against the owning integration. - integrationID, err := s.q.GetInteractionCallbackIntegrationID(ctx, dbsqlc.GetInteractionCallbackIntegrationIDParams{ - ID: id, - }) + integrationID, err := s.q.GetInteractionCallbackIntegrationID( + ctx, + dbsqlc.GetInteractionCallbackIntegrationIDParams{ + ID: id, + }, + ) if errors.Is(err, pgx.ErrNoRows) { return uuid.Nil, storeerr.ErrNotFound } @@ -620,14 +556,15 @@ func (s *Store) GetInteractionForHandlerCallback( func (s *Store) ResolveAgentInteractionFromHandler( ctx context.Context, input ResolveAgentInteractionFromHandlerInput, ) (AgentInteractionRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentInteractionRecord{}, err } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() _, destination, err := lockAuthorizedInteractionDestination( ctx, - tx, + unit, input.ProjectID, input.AgentID, input.ID, @@ -652,23 +589,15 @@ func (s *Store) ResolveAgentInteractionFromHandler( if err := validateIntegrationInputActor(integration, input.Actor); err != nil { return AgentInteractionRecord{}, err } - notifications := s.newTxNotifications() - record, err := resolveAgentInteractionTx( + record, err := resolveExecutionInteraction( ctx, - notifications, - tx, - dbsqlc.New(tx), + unit, input.ResolveAgentInteractionInput, ) if err != nil { return AgentInteractionRecord{}, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - notifications, - "resolve interaction from handler", - ); err != nil { + if err := unit.Commit(ctx, "resolve interaction from handler"); err != nil { return AgentInteractionRecord{}, err } return record, nil @@ -677,18 +606,19 @@ func (s *Store) ResolveAgentInteractionFromHandler( func (s *Store) GetAgentInteractionForPresentation( ctx context.Context, projectID, agentID, id uuid.UUID, ) (AgentInteractionRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentInteractionRecord{}, err } - defer func() { _ = tx.Rollback(ctx) }() - record, _, err := lockAuthorizedInteractionDestination(ctx, tx, projectID, agentID, id) + defer func() { _ = unit.Rollback(ctx) }() + record, _, err := lockAuthorizedInteractionDestination(ctx, unit, projectID, agentID, id) return record, err } func lockAuthorizedInteractionDestination( - ctx context.Context, tx pgx.Tx, projectID, agentID, id uuid.UUID, + ctx context.Context, unit *agentexecution.Unit, projectID, agentID, id uuid.UUID, ) (AgentInteractionRecord, *InteractionDestination, error) { + tx := unit.DB() q := dbsqlc.New(tx) params := dbsqlc.GetAgentInteractionParams{ProjectID: projectID, AgentID: agentID, ID: id} row, err := q.GetAgentInteraction(ctx, params) @@ -722,9 +652,9 @@ func lockAuthorizedInteractionDestination( ); err != nil { return AgentInteractionRecord{}, nil, err } - if _, err := q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: projectID, ID: agentID, - }); err != nil { + }, agentexecution.ExternalAuthority{}); err != nil { return AgentInteractionRecord{}, nil, err } _, contract, err := loadInteractionContract(ctx, q, projectID, agentID) @@ -761,15 +691,16 @@ func (s *Store) RecordInteractionPresentationReceipt( if err := validateInteractionObject(input.Receipt, InteractionReceiptMaxBytes); err != nil { return AgentInteractionRecord{}, storeerr.InvalidRequest(err) } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return AgentInteractionRecord{}, err } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() q := dbsqlc.New(tx) - if _, err := q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: input.ProjectID, ID: input.AgentID, - }); err != nil { + }, agentexecution.ExternalAuthority{}); err != nil { if errors.Is(err, pgx.ErrNoRows) { return AgentInteractionRecord{}, storeerr.ErrNotFound } @@ -801,26 +732,81 @@ func (s *Store) RecordInteractionPresentationReceipt( } return record, nil } - receiptParams := dbsqlc.RecordAgentInteractionPresentationReceiptParams{ - ProjectID: input.ProjectID, AgentID: input.AgentID, ID: input.ID, - Destination: record.Destination, Receipt: input.Receipt, - } - changed, err := q.RecordAgentInteractionPresentationReceipt(ctx, receiptParams) + h, err := unit.Handle(input.ProjectID, input.AgentID) if err != nil { - return AgentInteractionRecord{}, fmt.Errorf( - "record interaction presentation receipt: %w", - err, - ) + return AgentInteractionRecord{}, err } - if changed != 1 { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict + if _, err = h.RecordPresentation(ctx, input.ID, record.Destination, input.Receipt); err != nil { + return AgentInteractionRecord{}, err } row, err = q.GetAgentInteraction(ctx, params) if err != nil { return AgentInteractionRecord{}, err } - if err := tx.Commit(ctx); err != nil { + if err := unit.Commit(ctx, "RecordInteractionPresentationReceipt"); err != nil { return AgentInteractionRecord{}, err } return agentInteractionRecordFromSQLC(row), nil } + +func interactionSelectionForOrigin( + ctx context.Context, + tx dbsqlc.DBTX, + projectID, agentID, configID, originTargetID uuid.UUID, +) (InteractionSelection, error) { + q := dbsqlc.New(tx) + selection := InteractionSelection{AutoSelect: true} + if originTargetID == uuid.Nil { + return selection, nil + } + contract, err := interactionContractForConfig(ctx, q, projectID, configID) + if err != nil { + return InteractionSelection{}, err + } + handlers, err := prepareInteractionHandlers(ctx, q, projectID, contract.InteractionHandlers) + if err != nil { + return InteractionSelection{}, err + } + if len(handlers) == 0 { + return selection, nil + } + target, err := q.GetInteractionDestinationTarget( + ctx, + dbsqlc.GetInteractionDestinationTargetParams{ + ProjectID: projectID, + AgentID: agentID, + TargetID: originTargetID, + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return selection, nil + } + if err != nil { + return InteractionSelection{}, err + } + matches := 0 + for key, handler := range handlers { + if target.IntegrationState != string(integrationstore.IntegrationStateActive) || + handler.prepared.IntegrationID != target.IntegrationID { + continue + } + destination, err := resolveInteractionHandlerDestination(ctx, tx, projectID, agentID, key, handler) + if err != nil { + return InteractionSelection{}, err + } + if destination == nil || destination.IntegrationTargetID != originTargetID { + continue + } + + matches++ + selection = InteractionSelection{ + AutoSelect: true, + IntegrationTargetID: originTargetID, + HandlerKey: key, + } + } + if matches != 1 { + selection = InteractionSelection{AutoSelect: true} + } + return selection, nil +} diff --git a/internal/storage/executionstore/integration_testmain_test.go b/internal/storage/executionstore/integration_testmain_test.go index 4be6c3d30..51a4d2096 100644 --- a/internal/storage/executionstore/integration_testmain_test.go +++ b/internal/storage/executionstore/integration_testmain_test.go @@ -151,14 +151,6 @@ func normalizedJSON(value json.RawMessage) json.RawMessage { return value } -func marshalJSON(value any) (json.RawMessage, error) { - body, err := json.Marshal(value) - if err != nil { - return nil, err - } - return body, nil -} - func sameJSON(a, b json.RawMessage) bool { return jsoncanonical.Equal(a, b) } diff --git a/internal/storage/executionstore/interaction_selection_admission.go b/internal/storage/executionstore/interaction_selection_admission.go index ae327477d..ab15183c6 100644 --- a/internal/storage/executionstore/interaction_selection_admission.go +++ b/internal/storage/executionstore/interaction_selection_admission.go @@ -5,17 +5,30 @@ import ( "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) -func selectAdmittedInteractionDestinationTx( +func applyAdmissionDestination( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, projectID, agentID uuid.UUID, - inputs []AgentInputRecord, + admission agentexecution.InputAdmission, ) error { + if len(admission.Inputs) == 0 { + return nil + } + tx := unit.DB() + inputs := make([]AgentInputRecord, len(admission.Inputs)) + for i, input := range admission.Inputs { + inputs[i] = AgentInputRecord{ + ID: input.ID, + ActorID: input.ActorID, + IntegrationTargetID: input.IntegrationTargetID, + InputKind: "content", + } + } q := dbsqlc.New(tx) row, err := q.GetInteractionSelection(ctx, dbsqlc.GetInteractionSelectionParams{ ProjectID: projectID, AgentID: agentID, @@ -59,21 +72,35 @@ func selectAdmittedInteractionDestinationTx( return fmt.Errorf("load admitted input actor identities: %w", err) } for _, actor := range actors { - preserves[actor.ID] = internalInputPreservesInteractionSelection(ActorProvider(actor.Provider), actor.ProviderUserID) + preserves[actor.ID] = internalInputPreservesInteractionSelection( + ActorProvider(actor.Provider), + actor.ProviderUserID, + ) } } for i := len(inputs) - 1; i >= 0; i-- { input := inputs[i] - if input.InputKind != "content" || (input.IntegrationTargetID == uuid.Nil && preserves[input.ActorID]) { + if input.InputKind != "content" || + (input.IntegrationTargetID == uuid.Nil && preserves[input.ActorID]) { continue } - selection, err := resolveInteractionSelectionForOriginTx( + selection, err := interactionSelectionForOrigin( ctx, tx, projectID, agentID, row.CurrentConfigID, input.IntegrationTargetID, ) if err != nil || selection == interactionSelectionFromRow(row) { return err } - return writeInteractionSelection(ctx, q, projectID, agentID, selection) + h, err := unit.Handle(projectID, agentID) + if err != nil { + return err + } + _, err = h.SetInteractionTarget( + ctx, + selection.IntegrationTargetID, + selection.HandlerKey, + selection.AutoSelect, + ) + return err } return nil } diff --git a/internal/storage/executionstore/interaction_selection_export_integration_test.go b/internal/storage/executionstore/interaction_selection_export_integration_test.go index effea71af..cb5841eb4 100644 --- a/internal/storage/executionstore/interaction_selection_export_integration_test.go +++ b/internal/storage/executionstore/interaction_selection_export_integration_test.go @@ -8,26 +8,45 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) -func IntegrationSelectAdmittedInteractionDestinationTx( - ctx context.Context, tx pgx.Tx, projectID, agentID uuid.UUID, inputs []AgentInputRecord, +func IntegrationApplyAdmissionDestination( + ctx context.Context, + unit *agentexecution.Unit, + projectID, agentID uuid.UUID, + inputs []AgentInputRecord, ) error { - return selectAdmittedInteractionDestinationTx(ctx, tx, projectID, agentID, inputs) + admission := agentexecution.InputAdmission{} + for _, input := range inputs { + admission.Inputs = append( + admission.Inputs, + agentexecution.AdmittedContent{ + ID: input.ID, + ActorID: input.ActorID, + IntegrationTargetID: input.IntegrationTargetID, + }, + ) + } + return applyAdmissionDestination(ctx, unit, projectID, agentID, admission) } func IntegrationSelectInteractionDestinationForOriginTx( - ctx context.Context, tx pgx.Tx, projectID, agentID, originTargetID uuid.UUID, + ctx context.Context, + unit *agentexecution.Unit, + projectID, agentID, originTargetID uuid.UUID, ) (InteractionSelection, error) { - if err := selectAdmittedInteractionDestinationTx(ctx, tx, projectID, agentID, []AgentInputRecord{{ - InputKind: "content", IntegrationTargetID: originTargetID, - }}); err != nil { + if err := IntegrationApplyAdmissionDestination(ctx, + unit, + projectID, + agentID, + []AgentInputRecord{{InputKind: "content", + IntegrationTargetID: originTargetID}}); err != nil { return InteractionSelection{}, err } - row, err := dbsqlc.New(tx).GetInteractionSelection(ctx, dbsqlc.GetInteractionSelectionParams{ - ProjectID: projectID, AgentID: agentID, - }) + row, err := dbsqlc.New(unit.DB()). + GetInteractionSelection(ctx, dbsqlc.GetInteractionSelectionParams{ProjectID: projectID, AgentID: agentID}) return interactionSelectionFromRow(row), err } diff --git a/internal/storage/executionstore/interaction_selection_producers_integration_test.go b/internal/storage/executionstore/interaction_selection_producers_integration_test.go index 89140f2cf..8e20feb96 100644 --- a/internal/storage/executionstore/interaction_selection_producers_integration_test.go +++ b/internal/storage/executionstore/interaction_selection_producers_integration_test.go @@ -36,14 +36,15 @@ func TestSubagentCompletedReportPreservesInteractionSelection(t *testing.T) { ) require.True(t, found) require.Len(t, opening.Inputs, 1) - claim, err := f.store.Execution().ClaimNormalModelCall(f.ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: child.Agent.ID, - RuntimeLockID: childLock.ID, - OpeningInputIDs: []uuid.UUID{opening.Inputs[0].ID}, - AgentConfigID: child.AgentConfig.ID, - InputEventSequence: opening.Events[len(opening.Events)-1].Sequence, - }) + prepared1, err := f.store.Execution(). + PrepareNormalModelCall(f.ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: child.Agent.ID, + RuntimeLockID: childLock.ID, + OpeningInputIDs: []uuid.UUID{opening.Inputs[0].ID}, + }) + claim := prepared1.Claim + require.NoError(t, err) providerModelSlug := modelProviderSlugForContext( t, f.ctx, f.store, testProjectID, child.Agent.ID, claim.Context.ID, @@ -86,7 +87,12 @@ func TestSubagentCompletedReportPreservesInteractionSelection(t *testing.T) { require.Contains(t, string(turnEvents[0].ContentBlocks), report) selection, err := f.store.Execution().GetInteractionSelection(f.ctx, testProjectID, parent.ID) require.NoError(t, err) - require.Equal(t, before, selection, "admitting the completed subagent report must preserve the Slack selection") + require.Equal( + t, + before, + selection, + "admitting the completed subagent report must preserve the Slack selection", + ) } func TestCronTriggerAgentInputPreservesInteractionSelection(t *testing.T) { @@ -99,7 +105,8 @@ func TestCronTriggerAgentInputPreservesInteractionSelection(t *testing.T) { input := cronTriggerInput("Scheduled review", f.process.AgentID, true) trigger, err := f.store.Execution().CreateCronTrigger(f.ctx, input) require.NoError(t, err) - _, err = f.store.pool.Exec(f.ctx, + _, err = f.store.pool.Exec( + f.ctx, `UPDATE cron_triggers SET next_fire_after = statement_timestamp() - interval '1 minute' WHERE id = $1`, trigger.ID, ) diff --git a/internal/storage/executionstore/interaction_selection_queries_integration_test.go b/internal/storage/executionstore/interaction_selection_queries_integration_test.go index 03518afd9..56f8cad15 100644 --- a/internal/storage/executionstore/interaction_selection_queries_integration_test.go +++ b/internal/storage/executionstore/interaction_selection_queries_integration_test.go @@ -3,143 +3,57 @@ package executionstore_test import ( - "context" "encoding/json" - "strings" "testing" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgconn" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/testutil/integrationdb" "github.com/stretchr/testify/require" ) -type interactionQueryTx struct { - pgx.Tx - queries []string -} - -func (tx *interactionQueryTx) record(sql string) { - line, _, _ := strings.Cut(sql, "\n") - tx.queries = append(tx.queries, strings.Fields(line)[2]) -} - -func (tx *interactionQueryTx) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { - tx.record(sql) - return tx.Tx.QueryRow(ctx, sql, args...) -} - -func (tx *interactionQueryTx) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { - tx.record(sql) - return tx.Tx.Query(ctx, sql, args...) //nolint:rowserrcheck // SQLC consumes the returned rows and checks Err. -} - -func (tx *interactionQueryTx) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { - tx.record(sql) - return tx.Tx.Exec(ctx, sql, args...) -} - -func TestInteractionSelectionConditionalQueries(t *testing.T) { - t.Parallel() - f := newIntegrationInteractionFixture(t) - actor, err := executionstore.CronTriggerActor(testOrgID, uuid.New(), "Routine instructions") - require.NoError(t, err) - input, _, _, err := f.store.Execution().CreateAgentContentInput(f.ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, AgentID: f.process.AgentID, Actor: actor, - ContentBlocks: json.RawMessage(`[{"type":"text","text":"routine"}]`), - }) - require.NoError(t, err) - f.change(t, nil) - tx := &interactionQueryTx{Tx: integrationdb.BeginTx(t, f.ctx, f.store.pool)} - q := dbsqlc.New(tx) - _, err = q.LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: testProjectID, ID: f.process.AgentID, - }) - require.NoError(t, err) - tx.queries = nil - require.NoError(t, executionstore.IntegrationSelectAdmittedInteractionDestinationTx( - f.ctx, tx, testProjectID, f.process.AgentID, []executionstore.AgentInputRecord{input}, - )) - require.Equal(t, []string{"GetInteractionSelection"}, tx.queries, "ordinary admission needs only the agent row") - tx.queries = nil - destination, err := executionstore.IntegrationCaptureInteractionDestinationTx( - f.ctx, tx, testProjectID, f.process.AgentID, - ) - require.NoError(t, err) - require.Nil(t, destination) - require.Equal(t, []string{"GetInteractionSelection"}, tx.queries, "dashboard prompt capture skips config resolution") - tx.queries = nil - input.IntegrationTargetID = f.a.ID - require.NoError(t, executionstore.IntegrationSelectAdmittedInteractionDestinationTx( - f.ctx, tx, testProjectID, f.process.AgentID, []executionstore.AgentInputRecord{input}, - )) - require.Equal(t, []string{"GetInteractionSelection", "GetAgentConfig"}, tx.queries, - "an origin without configured handlers needs no integration or actor lookup") - _, err = q.SetInteractionSelection(f.ctx, dbsqlc.SetInteractionSelectionParams{ - ProjectID: testProjectID, AgentID: f.process.AgentID, AutoSelect: false, - TargetID: &f.a.ID, HandlerKey: new("chat"), - }) - require.NoError(t, err) - tx.queries = nil - require.NoError(t, executionstore.IntegrationSelectAdmittedInteractionDestinationTx( - f.ctx, tx, testProjectID, f.process.AgentID, []executionstore.AgentInputRecord{input}, - )) - require.Equal(t, []string{"GetInteractionSelection"}, tx.queries, "manual pins bypass origin resolution") - tx.queries = nil - selection, err := f.store.Execution().ReconcileInteractionSelectionTx(f.ctx, tx, testProjectID, f.process.AgentID) - require.NoError(t, err) - require.Equal(t, executionstore.InteractionSelection{AutoSelect: false}, selection, - "a removed handler clears the destination while preserving manual mode") - require.Equal(t, []string{"GetInteractionSelection", "GetAgentConfig", "SetInteractionSelection"}, tx.queries) -} - func TestInteractionSelectionOriginPrecedesBackgroundActor(t *testing.T) { t.Parallel() f := newIntegrationInteractionFixture(t) actor, err := executionstore.CronTriggerActor(testOrgID, uuid.New(), "Scheduled thread") require.NoError(t, err) - input, _, _, err := f.store.Execution().CreateAgentContentInput(f.ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, AgentID: f.process.AgentID, Actor: actor, - ContentBlocks: json.RawMessage(`[{"type":"text","text":"scheduled"}]`), - }) + input, _, _, err := f.store.Execution(). + CreateAgentContentInput(f.ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, AgentID: f.process.AgentID, Actor: actor, + ContentBlocks: json.RawMessage(`[{"type":"text","text":"scheduled"}]`), + }) require.NoError(t, err) - tx := &interactionQueryTx{Tx: integrationdb.BeginTx(t, f.ctx, f.store.pool)} - q := dbsqlc.New(tx) - _, err = q.LockAgentInProject(f.ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: testProjectID, ID: f.process.AgentID, - }) + unit, err := f.store.Execution().IntegrationBeginUnit(f.ctx) require.NoError(t, err) + defer func() { _ = unit.Rollback(f.ctx) }() + _, err = unit.LockAgent( + f.ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: f.process.AgentID}, + agentexecution.ExternalAuthority{}, + ) + require.NoError(t, err) + q := dbsqlc.New(unit.DB()) background := input input.IntegrationTargetID = f.b.ID - tx.queries = nil - require.NoError(t, executionstore.IntegrationSelectAdmittedInteractionDestinationTx( - f.ctx, tx, testProjectID, f.process.AgentID, []executionstore.AgentInputRecord{input, background, background}, + require.NoError(t, executionstore.IntegrationApplyAdmissionDestination( + f.ctx, + unit, + testProjectID, + f.process.AgentID, + []executionstore.AgentInputRecord{input, background, background}, )) - require.Equal(t, 1, countInteractionQuery(tx.queries, "ListActorIdentitiesByIDs")) - require.Equal(t, 1, countInteractionQuery(tx.queries, "GetInteractionSelection")) selection, err := q.GetInteractionSelection(f.ctx, dbsqlc.GetInteractionSelectionParams{ ProjectID: testProjectID, AgentID: f.process.AgentID, }) require.NoError(t, err) require.Equal(t, "other", selection.HandlerKey) require.Equal(t, f.b.ID, *selection.InteractionTargetID) - tx.queries = nil - require.NoError(t, executionstore.IntegrationSelectAdmittedInteractionDestinationTx( - f.ctx, tx, testProjectID, f.process.AgentID, []executionstore.AgentInputRecord{background, background}, + require.NoError(t, executionstore.IntegrationApplyAdmissionDestination( + f.ctx, + unit, + testProjectID, + f.process.AgentID, + []executionstore.AgentInputRecord{background, background}, )) - require.Equal(t, []string{"GetInteractionSelection", "ListActorIdentitiesByIDs"}, tx.queries, - "background content preserves a selected handler without loading config or integrations") -} - -func countInteractionQuery(queries []string, name string) int { - count := 0 - for _, query := range queries { - if query == name { - count++ - } - } - return count } diff --git a/internal/storage/executionstore/json_helpers.go b/internal/storage/executionstore/json_helpers.go index 934f0923f..4cb058cbb 100644 --- a/internal/storage/executionstore/json_helpers.go +++ b/internal/storage/executionstore/json_helpers.go @@ -22,18 +22,6 @@ func normalizedJSONArray(value json.RawMessage) json.RawMessage { return value } -func normalizedJSONObject(value json.RawMessage, fieldName string) (json.RawMessage, error) { - value = normalizedJSON(value) - var object map[string]json.RawMessage - if err := json.Unmarshal(value, &object); err != nil { - return nil, fmt.Errorf("parse %s: %w", fieldName, err) - } - if object == nil { - return nil, fmt.Errorf("%s must be a JSON object", fieldName) - } - return value, nil -} - func marshalJSON(value any) (json.RawMessage, error) { body, err := json.Marshal(value) if err != nil { diff --git a/internal/storage/executionstore/kernel_event_authority_checkpoints_integration_test.go b/internal/storage/executionstore/kernel_event_authority_checkpoints_integration_test.go index 69b9b4619..577ec9b18 100644 --- a/internal/storage/executionstore/kernel_event_authority_checkpoints_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_checkpoints_integration_test.go @@ -246,7 +246,11 @@ func TestContextCheckpointPublicationRejectsRepeatedTransition(t *testing.T) { ctx, input, ); !errors.Is(err, storeerr.ErrStateTransitionConflict) { - t.Fatalf("repeat checkpoint publication error = %v, want %v", err, storeerr.ErrStateTransitionConflict) + t.Fatalf( + "repeat checkpoint publication error = %v, want %v", + err, + storeerr.ErrStateTransitionConflict, + ) } } @@ -307,7 +311,11 @@ func TestContextCheckpointLineageDerivesPriorCheckpointFromEvents(t *testing.T) if err != nil { t.Fatalf("publish second checkpoint: %v", err) } - secondSourceStart, err := executionstore.IntegrationCompactionSourceStartTx(ctx, fixture.Store.q, secondClaim.Context) + secondSourceStart, err := executionstore.IntegrationCompactionSourceStartTx( + ctx, + fixture.Store.pool, + secondClaim.Context, + ) if err != nil { t.Fatalf("derive second compaction source start: %v", err) } @@ -316,7 +324,11 @@ func TestContextCheckpointLineageDerivesPriorCheckpointFromEvents(t *testing.T) } if secondCheckpoint.SummarizedThroughEventSequence != secondFrontier || secondCheckpoint.SummarizedThroughEventSequence <= firstCheckpoint.SummarizedThroughEventSequence { - t.Fatalf("checkpoint lineage did not advance: first=%+v second=%+v", firstCheckpoint, secondCheckpoint) + t.Fatalf( + "checkpoint lineage did not advance: first=%+v second=%+v", + firstCheckpoint, + secondCheckpoint, + ) } } @@ -524,7 +536,13 @@ func TestKernelContextCheckpointRejectsOpenToolCallAuthority(t *testing.T) { t.Run(state, func(t *testing.T) { fixture := newProcessDaemonFixture(t, ctx, "kernel_checkpoint_"+state+"_tool_call") now := fixture.Now.Add(time.Minute) - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "kernel_checkpoint_"+state+"_tool_call", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "kernel_checkpoint_"+state+"_tool_call", + "read_process", + ) if state == "running" { claimToolCallForTest( t, @@ -537,9 +555,24 @@ func TestKernelContextCheckpointRejectsOpenToolCallAuthority(t *testing.T) { ) } sourceSequence := toolCallSourceSequenceForCheckpointTest(t, ctx, fixture, toolCallID) - claim := claimSentCompactionForRangeTest(t, ctx, fixture, 1, sourceSequence, sourceSequence, now.Add(time.Second)) - _, err := publishCheckpointForRangeTest(t, ctx, fixture, claim, "unsafe-"+state, now.Add(2*time.Second)) - if err == nil || !strings.Contains(err.Error(), "cuts an active tool call") { + claim := claimSentCompactionForRangeTest( + t, + ctx, + fixture, + 1, + sourceSequence, + sourceSequence, + now.Add(time.Second), + ) + _, err := publishCheckpointForRangeTest( + t, + ctx, + fixture, + claim, + "unsafe-"+state, + now.Add(2*time.Second), + ) + if !errors.Is(err, storeerr.ErrCheckpointBoundaryUnsafe) { t.Fatalf("checkpoint %s tool authority error = %v", state, err) } }) @@ -548,12 +581,19 @@ func TestKernelContextCheckpointRejectsOpenToolCallAuthority(t *testing.T) { t.Run("completed", func(t *testing.T) { fixture := newProcessDaemonFixture(t, ctx, "kernel_checkpoint_completed_tool_call") now := fixture.Now.Add(time.Minute) - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "kernel_checkpoint_completed_tool_call", "read_process") - completed, err := fixture.Store.Execution().CompleteToolCall(ctx, executionstore.CompleteToolCallInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, ID: toolCallID, - Outcome: executionstore.ToolResultOutcomeSucceeded, RuntimeLockID: fixture.Lock.ID, - ResultContentParts: json.RawMessage(`[{"type":"structured_data","value":{"ok":true}}]`), - }) + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "kernel_checkpoint_completed_tool_call", + "read_process", + ) + completed, err := fixture.Store.Execution(). + CompleteToolCall(ctx, executionstore.CompleteToolCallInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, ID: toolCallID, + Outcome: executionstore.ToolResultOutcomeSucceeded, RuntimeLockID: fixture.Lock.ID, + ResultContentParts: json.RawMessage(`[{"type":"structured_data","value":{"ok":true}}]`), + }) if err != nil { t.Fatalf("complete tool call: %v", err) } @@ -564,12 +604,23 @@ func TestKernelContextCheckpointRejectsOpenToolCallAuthority(t *testing.T) { claim := claimSentCompactionForRangeTest( t, ctx, fixture, 1, appended[0].Sequence, appended[0].Sequence, now.Add(time.Second), ) - checkpoint, err := publishCheckpointForRangeTest(t, ctx, fixture, claim, "safe", now.Add(2*time.Second)) + checkpoint, err := publishCheckpointForRangeTest( + t, + ctx, + fixture, + claim, + "safe", + now.Add(2*time.Second), + ) if err != nil { t.Fatalf("checkpoint closed tool authority: %v", err) } if checkpoint.SummarizedThroughEventSequence != appended[0].Sequence { - t.Fatalf("checkpoint end = %d, want %d", checkpoint.SummarizedThroughEventSequence, appended[0].Sequence) + t.Fatalf( + "checkpoint end = %d, want %d", + checkpoint.SummarizedThroughEventSequence, + appended[0].Sequence, + ) } }) } diff --git a/internal/storage/executionstore/kernel_event_authority_constraints_integration_test.go b/internal/storage/executionstore/kernel_event_authority_constraints_integration_test.go index 8bf91434d..5242d011f 100644 --- a/internal/storage/executionstore/kernel_event_authority_constraints_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_constraints_integration_test.go @@ -5,7 +5,6 @@ package executionstore_test import ( "context" "encoding/json" - "strings" "testing" "time" @@ -192,7 +191,11 @@ WHERE agent_id = $1 AND id = $2 t.Fatalf("transition tool call to %s: %v", state, err) } if result.RowsAffected() != 1 { - t.Fatalf("transition tool call to %s affected %d rows, want 1", state, result.RowsAffected()) + t.Fatalf( + "transition tool call to %s affected %d rows, want 1", + state, + result.RowsAffected(), + ) } } for _, state := range paths[transition.from] { @@ -243,7 +246,7 @@ VALUES ($1, $2, 'graph_invalid_initial_state', 'read_process', '{}', 'built_in', var incomplete bool if err := fixture.Store.pool.QueryRow( ctx, - `SELECT agent_has_incomplete_tool_batch($1, $2)`, + `SELECT h.incomplete_tools FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND a.id=$2`, testProjectID, fixture.AgentID, ).Scan(&incomplete); err != nil || !incomplete { @@ -280,7 +283,7 @@ WHERE agent_id = $1 AND id = $2 } if err := fixture.Store.pool.QueryRow( ctx, - `SELECT agent_has_incomplete_tool_batch($1, $2)`, + `SELECT h.incomplete_tools FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND a.id=$2`, testProjectID, fixture.AgentID, ).Scan(&incomplete); err != nil || incomplete { @@ -336,7 +339,11 @@ WHERE project_id = $1 AND id = $5 `, testProjectID, turnID, "agent_input:"+inputID.String(), inputID, fixture.AgentID); err != nil { t.Fatalf("insert orphan agent-input event: %v", err) } - if _, err := tx.Exec(ctx, `SET CONSTRAINTS agent_input_events_admission_valid IMMEDIATE`); !isPgCheckViolation(err) { + if _, + err := tx.Exec(ctx, + `SET CONSTRAINTS agent_input_events_admission_valid IMMEDIATE`); !isPgCheckViolation( + err, + ) { t.Fatalf("orphan agent-input event error = %v, want check violation", err) } } @@ -412,7 +419,13 @@ func TestKernelTurnOpeningEventMustBeAgentInput(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "kernel_turn_opening_event_kind") - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "kernel_turn_opening_event_kind", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "kernel_turn_opening_event_kind", + "read_process", + ) turnID := turnIDForProcessToolCallTest(t, ctx, fixture, toolCallID) contextID := modelOutputContextForTurnTest( t, @@ -561,48 +574,14 @@ func TestKernelContentBlockOwnerAndOrdinalConstraints(t *testing.T) { ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "kernel_content_block_owner_ordinal") now := fixture.Now.Add(time.Minute) - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "kernel_content_block_owner_ordinal", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "kernel_content_block_owner_ordinal", + "read_process", + ) modelOutputID := modelOutputIDForToolCall(t, ctx, fixture.Store, fixture.AgentID, toolCallID) - for index, test := range []struct { - name string - input executionstore.CreateContentBlockInput - }{ - { - name: "text", - input: executionstore.CreateContentBlockInput{ - BlockKind: executionstore.ContentBlockKindText, - TextContent: "before\x00after", - }, - }, - { - name: "structured data", - input: executionstore.CreateContentBlockInput{ - BlockKind: executionstore.ContentBlockKindStructuredData, - StructuredData: json.RawMessage(`{"value":"before\u0000after"}`), - }, - }, - { - name: "metadata", - input: executionstore.CreateContentBlockInput{ - BlockKind: executionstore.ContentBlockKindText, - TextContent: "safe", - Metadata: map[string]string{"key": "before\x00after"}, - }, - }, - } { - t.Run("database unsafe "+test.name, func(t *testing.T) { - t.Parallel() - test.input.ProjectID = testProjectID - test.input.AgentID = fixture.AgentID - test.input.OwnerKind = executionstore.ContentBlockOwnerModelOutput - test.input.OwnerModelOutputID = modelOutputID - test.input.Ordinal = int32(index + 10) - _, err := executionstore.IntegrationCreateContentBlockTx(ctx, fixture.Store.pool, test.input) - if err == nil || !strings.Contains(err.Error(), "U+0000") { - t.Fatalf("create database-unsafe content block error = %v, want U+0000", err) - } - }) - } var block executionstore.ContentBlockRecord var blockKind string @@ -751,7 +730,10 @@ WHERE agent_id = $1 AND id = $2 control_type = 'cancel_current' WHERE project_id = $1 AND agent_id = $2 AND id = $3 `, testProjectID, fixture.AgentID, contentInputID); !isPgCode(err, "25006") { - t.Fatalf("mutating content-bearing input subtype error = %v, want immutable-event SQLSTATE 25006", err) + t.Fatalf( + "mutating content-bearing input subtype error = %v, want immutable-event SQLSTATE 25006", + err, + ) } } @@ -867,15 +849,18 @@ func TestModelOutputContentBlockInsertSerializesWithContextClosure(t *testing.T) if err != nil { t.Fatalf("create model output for owner lock: %v", err) } - if _, err := executionstore.IntegrationCreateContentBlockTx(ctx, blockTx, executionstore.CreateContentBlockInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - OwnerKind: executionstore.ContentBlockOwnerModelOutput, - OwnerModelOutputID: output.ID, - Ordinal: 0, - BlockKind: executionstore.ContentBlockKindText, - TextContent: "concurrent model output block", - }); err != nil { + if _, + err := executionstore.IntegrationCreateContentBlockTx(ctx, + blockTx, + executionstore.CreateContentBlockInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + OwnerKind: executionstore.ContentBlockOwnerModelOutput, + OwnerModelOutputID: output.ID, + Ordinal: 0, + BlockKind: executionstore.ContentBlockKindText, + TextContent: "concurrent model output block", + }); err != nil { t.Fatalf("create model output block for owner lock: %v", err) } @@ -908,98 +893,6 @@ WHERE project_id = $1 AND agent_id = $2 AND id = $3 AND state = 'started' } } -func TestKernelTypedFrontierAddAgentEventOnlyOnRealInsert(t *testing.T) { - t.Parallel() - ctx := context.Background() - pool := openIntegrationDB(t, ctx) - seedMigratedDB(t, ctx, pool) - store := newIntegrationStore(pool) - agentID := mustCreateAgent(t, ctx, store) - user := mustCreateProjectOperatorUser(t, ctx, store, "wakeup-idem@example.com", "Wakeup Idempotency") - input, _, _, err := store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: agentID, - Actor: mustOmnaraActorParams(t, user.ID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"wakeup test"}]`), - IdempotencyKey: "wakeup-idem-input", - }) - if err != nil { - t.Fatalf("create input: %v", err) - } - - var turnID uuid.UUID - if err := pool.QueryRow( - ctx, - `SELECT event.turn_id FROM agent_events event JOIN agents agent ON agent.id = event.agent_id WHERE agent.project_id = $1 AND event.agent_id = $2 ORDER BY event.sequence DESC LIMIT 1`, - testProjectID, - agentID, - ).Scan(&turnID); err != nil { - t.Fatalf("load existing turn id: %v", err) - } - - countAgentEvents := func(txNotifications *notifications.TxNotifications) int { - recorder := &recordingPostCommitPublisher{} - txNotifications.Flush(ctx, recorder) - count := 0 - for _, intent := range recorder.intents { - event, ok := intent.(notifications.AgentEventCommitted) - if !ok { - continue - } - if event.AgentID == agentID { - count++ - } - } - return count - } - - txNotifications := notifications.NewTxNotifications() - tx, err := pool.Begin(ctx) - if err != nil { - t.Fatalf("begin tx: %v", err) - } - defer func() { _ = tx.Rollback(ctx) }() - - if _, err := executionstore.IntegrationAppendTypedAgentEventTx( - ctx, - txNotifications, - tx, - executionstore.AppendTypedAgentEventInput{ - ProjectID: testProjectID, - AgentID: agentID, - TurnID: turnID, - Kind: events.KindAgentInput, - IdempotencyKey: "agent_input:" + input.ID.String(), - AgentInputID: input.ID, - }, - ); err != nil { - t.Fatalf("first append: %v", err) - } - if got := countAgentEvents(txNotifications); got != 1 { - t.Fatalf("AgentEventCommitted intents after real INSERT = %d, want 1", got) - } - - txn2 := notifications.NewTxNotifications() - if _, err := executionstore.IntegrationAppendTypedAgentEventTx( - ctx, - txn2, - tx, - executionstore.AppendTypedAgentEventInput{ - ProjectID: testProjectID, - AgentID: agentID, - TurnID: turnID, - Kind: events.KindAgentInput, - IdempotencyKey: "agent_input:" + input.ID.String(), - AgentInputID: input.ID, - }, - ); err != nil { - t.Fatalf("second append (pointer re-find): %v", err) - } - if got := countAgentEvents(txn2); got != 0 { - t.Fatalf("AgentEventCommitted intents on idempotent pointer re-find = %d, want 0", got) - } -} - func TestKernelTypedFrontierRejectsMismatchedPointer(t *testing.T) { t.Parallel() ctx := context.Background() @@ -1026,7 +919,7 @@ func TestKernelTypedFrontierRejectsMismatchedPointer(t *testing.T) { AgentInputID: testID("not_a_model_output"), }, ) - if err == nil || err.Error() != "model_output event requires model_output_id" { + if !isPgCheckViolation(err) { t.Fatalf("mismatched typed event error = %v, want model output validation error", err) } } @@ -1039,13 +932,14 @@ func TestKernelTypedFrontierRequiresTurnMembership(t *testing.T) { store := newIntegrationStore(pool) agentID := mustCreateAgent(t, ctx, store) user := mustCreateProjectOperatorUser(t, ctx, store, "orphan-event@example.com", "Orphan Event") - input, _, _, err := store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: agentID, - Actor: mustOmnaraActorParams(t, user.ID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"orphan"}]`), - IdempotencyKey: "orphan-event-input", - }) + input, _, _, err := store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: agentID, + Actor: mustOmnaraActorParams(t, user.ID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"orphan"}]`), + IdempotencyKey: "orphan-event-input", + }) if err != nil { t.Fatalf("create input for orphan event: %v", err) } diff --git a/internal/storage/executionstore/kernel_event_authority_continuation_integration_test.go b/internal/storage/executionstore/kernel_event_authority_continuation_integration_test.go index 9520031eb..d7a448c93 100644 --- a/internal/storage/executionstore/kernel_event_authority_continuation_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_continuation_integration_test.go @@ -25,7 +25,13 @@ func TestCompletedToolCallsForModelContextSpanTurnsAndRespectWatermark(t *testin ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "completed_tool_calls_at_watermark") now := fixture.Now.Add(time.Minute) - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "completed_tool_calls_at_watermark", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "completed_tool_calls_at_watermark", + "read_process", + ) modelContextID := modelContextIDForProcessToolCallTest(t, ctx, fixture, toolCallID) contextRecord, found, err := fixture.Store.Execution().GetModelCallContext( ctx, testProjectID, fixture.AgentID, modelContextID, @@ -185,9 +191,15 @@ func TestSemanticContextIdentityPreventsLaterOutputAtStaleFrontier(t *testing.T) t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "continuation_seed_later_unseen_result") - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "continuation_seed_later_unseen_result", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "continuation_seed_later_unseen_result", + "read_process", + ) modelContextID := modelContextIDForProcessToolCallTest(t, ctx, fixture, toolCallID) - initialContext, found, err := fixture.Store.Execution().GetModelCallContext( + _, found, err := fixture.Store.Execution().GetModelCallContext( ctx, testProjectID, fixture.AgentID, modelContextID, ) if err != nil { @@ -212,26 +224,29 @@ func TestSemanticContextIdentityPreventsLaterOutputAtStaleFrontier(t *testing.T) t.Fatalf("tool result events = %+v, want one", resultEvents) } openingInputID, _ := openingInputAndWatermarkForProcessToolCallTest(t, ctx, fixture, toolCallID) - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent: %v", err) } - _, err = fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{openingInputID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: initialContext.InputEventSequence, + _, err = fixture.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{openingInputID}, }) if !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("stale semantic model context claim error = %v, want %v", err, storeerr.ErrAgentNotAdvanceable) + t.Fatalf( + "stale semantic model context claim error = %v, want %v", + err, + storeerr.ErrAgentNotAdvanceable, + ) } seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("next continuation seed: %v", err) } - if !found || seed.TurnID != completed.TurnID || len(seed.InputIDs) != 1 || seed.InputIDs[0] != openingInputID { + if !found || seed.TurnID != completed.TurnID || len(seed.InputIDs) != 1 || + seed.InputIDs[0] != openingInputID { t.Fatalf( "continuation seed found=%v seed=%+v, want turn %s input %s", found, @@ -242,20 +257,35 @@ func TestSemanticContextIdentityPreventsLaterOutputAtStaleFrontier(t *testing.T) } } -func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { +func TestExecutionHeadExcludesHistoricalTurn(t *testing.T) { t.Parallel() ctx := context.Background() t.Run("active model context", func(t *testing.T) { t.Parallel() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "latest_frontier_active_context") + fixture, admitted, agent := newMultiInputContinuationSeedFixture( + t, + ctx, + "latest_frontier_active_context", + ) contextRecord := createContextForAdmittedTurnTest( - t, ctx, fixture, admitted, agent, "latest-frontier-active-context", fixture.Now.Add(4*time.Second), + t, + ctx, + fixture, + admitted, + agent, + "latest-frontier-active-context", + fixture.Now.Add(4*time.Second), ) assertFrontierCountTest( - t, ctx, fixture.Store, "active context before newer turn", - `SELECT count(*) FROM agent_continuable_model_contexts($1, $2) WHERE model_call_context_id = $3 AND NOT has_later_semantic_event`, - fixture.AgentID, contextRecord.ID, 1, + t, + ctx, + fixture.Store, + "active context before newer turn", + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN model_call_contexts c ON c.agent_id=h.agent_id AND c.id=h.normal_context_id WHERE a.project_id=$1 AND a.id=$2 AND c.id=$3 AND h.turn_continuable`, + fixture.AgentID, + contextRecord.ID, + 1, ) appendSyntheticLatestContentTurnForFrontierTest( @@ -263,15 +293,24 @@ func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { ) assertFrontierCountTest( - t, ctx, fixture.Store, "active context after newer turn", - `SELECT count(*) FROM agent_continuable_model_contexts($1, $2) WHERE model_call_context_id = $3 AND NOT has_later_semantic_event`, - fixture.AgentID, contextRecord.ID, 0, + t, + ctx, + fixture.Store, + "active context after newer turn", + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN model_call_contexts c ON c.agent_id=h.agent_id AND c.id=h.normal_context_id WHERE a.project_id=$1 AND a.id=$2 AND c.id=$3 AND h.turn_continuable`, + fixture.AgentID, + contextRecord.ID, + 0, ) }) t.Run("pending built-in tool", func(t *testing.T) { t.Parallel() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "latest_frontier_pending_tool") + fixture, admitted, agent := newMultiInputContinuationSeedFixture( + t, + ctx, + "latest_frontier_pending_tool", + ) contextRecord := createContextForAdmittedTurnTest( t, ctx, fixture, admitted, agent, "latest-frontier-pending-tool", fixture.Now.Add(4*time.Second), ) @@ -289,9 +328,14 @@ func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { fixture.Now.Add(5*time.Second), ) assertFrontierCountTest( - t, ctx, fixture.Store, "pending tool before newer turn", - `SELECT count(*) FROM agent_tool_work_frontiers($1, $2) WHERE model_call_context_id = $3`, fixture.AgentID, - contextRecord.ID, 1, + t, + ctx, + fixture.Store, + "pending tool before newer turn", + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN model_outputs o ON o.agent_id=h.agent_id AND o.id=h.pending_tool_output_id WHERE a.project_id=$1 AND a.id=$2 AND o.model_call_context_id=$3 AND h.incomplete_tools`, + fixture.AgentID, + contextRecord.ID, + 1, ) appendSyntheticLatestContentTurnForFrontierTest( @@ -299,26 +343,46 @@ func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { ) assertFrontierCountTest( - t, ctx, fixture.Store, "pending tool after newer turn", - `SELECT count(*) FROM agent_tool_work_frontiers($1, $2) WHERE model_call_context_id = $3`, fixture.AgentID, - contextRecord.ID, 0, + t, + ctx, + fixture.Store, + "pending tool after newer turn", + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN model_outputs o ON o.agent_id=h.agent_id AND o.id=h.pending_tool_output_id WHERE a.project_id=$1 AND a.id=$2 AND o.model_call_context_id=$3 AND h.incomplete_tools`, + fixture.AgentID, + contextRecord.ID, + 0, ) }) t.Run("completed tool result", func(t *testing.T) { t.Parallel() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "latest_frontier_completed_tool") + fixture, admitted, agent := newMultiInputContinuationSeedFixture( + t, + ctx, + "latest_frontier_completed_tool", + ) contextRecord := createContextForAdmittedTurnTest( - t, ctx, fixture, admitted, agent, "latest-frontier-completed-tool", fixture.Now.Add(4*time.Second), + t, + ctx, + fixture, + admitted, + agent, + "latest-frontier-completed-tool", + fixture.Now.Add(4*time.Second), ) completeToolCallForContinuationSeedTest( t, ctx, fixture, admitted.Turn.ID, contextRecord.ID, "latest_frontier_completed_tool", fixture.Now.Add(5*time.Second), ) assertFrontierCountTest( - t, ctx, fixture.Store, "completed tool before newer turn", - `SELECT count(*) FROM agent_model_result_frontiers($1, $2) WHERE model_call_context_id = $3`, - fixture.AgentID, contextRecord.ID, 1, + t, + ctx, + fixture.Store, + "completed tool before newer turn", + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN model_outputs o ON o.agent_id=h.agent_id AND o.id=h.pending_tool_output_id WHERE a.project_id=$1 AND a.id=$2 AND o.model_call_context_id=$3 AND NOT h.incomplete_tools`, + fixture.AgentID, + contextRecord.ID, + 1, ) appendSyntheticLatestContentTurnForFrontierTest( @@ -326,16 +390,27 @@ func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { ) assertFrontierCountTest( - t, ctx, fixture.Store, "completed tool after newer turn", - `SELECT count(*) FROM agent_model_result_frontiers($1, $2) WHERE model_call_context_id = $3`, - fixture.AgentID, contextRecord.ID, 0, + t, + ctx, + fixture.Store, + "completed tool after newer turn", + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN model_outputs o ON o.agent_id=h.agent_id AND o.id=h.pending_tool_output_id WHERE a.project_id=$1 AND a.id=$2 AND o.model_call_context_id=$3 AND NOT h.incomplete_tools`, + fixture.AgentID, + contextRecord.ID, + 0, ) }) t.Run("incomplete tool batch", func(t *testing.T) { t.Parallel() fixture := newProcessDaemonFixture(t, ctx, "latest_frontier_running_barrier") - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "latest_frontier_running_barrier", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "latest_frontier_running_barrier", + "read_process", + ) claimToolCallForTest( t, ctx, @@ -348,7 +423,7 @@ func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { var incomplete bool if err := fixture.Store.pool.QueryRow( ctx, - `SELECT agent_has_incomplete_tool_batch($1, $2)`, + `SELECT h.incomplete_tools FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND a.id=$2`, testProjectID, fixture.AgentID, ).Scan(&incomplete); err != nil || !incomplete { @@ -356,12 +431,16 @@ func TestLatestTurnFrontierHelpersIgnoreHistoricalTurn(t *testing.T) { } appendSyntheticLatestContentTurnForFrontierTest( - t, ctx, fixture, "latest_frontier_running_barrier_newer", fixture.Now.Add(time.Minute+time.Second), + t, + ctx, + fixture, + "latest_frontier_running_barrier_newer", + fixture.Now.Add(time.Minute+time.Second), ) if err := fixture.Store.pool.QueryRow( ctx, - `SELECT agent_has_incomplete_tool_batch($1, $2)`, + `SELECT h.incomplete_tools FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND a.id=$2`, testProjectID, fixture.AgentID, ).Scan(&incomplete); err != nil || incomplete { @@ -434,63 +513,17 @@ WHERE agent_id = $1 !strings.Contains(err.Error(), "completion state must match result existence") { t.Fatalf("complete sibling without typed result error = %v", err) } - if seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID); err != nil { + if seed, + found, + err := fixture.Store.Execution().NextAgentModelWork(ctx, + testProjectID, + fixture.AgentID); err != nil { t.Fatalf("next continuation seed: %v", err) } else if found { t.Fatalf("continuation seed = %+v, want none while sibling tool remains open", seed) } } -func TestContinuationSeedDoesNotResumeStoppedIncompleteBatch(t *testing.T) { - t.Parallel() - ctx := context.Background() - fixture, admitted, agent := newMultiInputContinuationSeedFixture( - t, - ctx, - "continuation_seed_stopped_sibling_result", - ) - now := fixture.Now.Add(time.Minute) - contextRecord := createContextForAdmittedTurnTest( - t, - ctx, - fixture, - admitted, - agent, - "continuation-seed-stopped-sibling-result", - now, - ) - _, toolCalls := recordToolCallBatchForContextTest( - t, - ctx, - fixture, - contextRecord.ID, - "continuation-seed-stopped-sibling-result", - []toolCallForContextTest{ - { - ProviderCallID: "call_continuation_seed_stopped_sibling_open", - Name: "read_process", - Input: json.RawMessage(`{}`), - }, - { - ProviderCallID: "call_continuation_seed_stopped_sibling_bad", - Name: "read_process", - Input: json.RawMessage(`{}`), - }, - }, - now.Add(time.Millisecond), - ) - if len(toolCalls) != 2 { - t.Fatalf("tool call batch = %+v, want two calls", toolCalls) - } - turnID := admitted.Turn.ID - appendCancelStopEventForContinuationSeedTest(t, ctx, fixture, turnID, now.Add(2*time.Second)) - if seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID); err != nil { - t.Fatalf("next continuation seed: %v", err) - } else if found { - t.Fatalf("continuation seed = %+v, want none after the turn was stopped", seed) - } -} - func TestContinuationSeedKeepsLaterToolResultAfterEarlierResultConsumed(t *testing.T) { t.Parallel() ctx := context.Background() @@ -503,7 +536,13 @@ func TestContinuationSeedKeepsLaterToolResultAfterEarlierResultConsumed(t *testi t, ctx, fixture, admitted.Turn.ID, firstContext.ID, "continuation_seed_later_result_first", fixture.Now.Add(5*time.Second), ) - firstResultEvents := listTypedToolResultEventsForToolCall(t, ctx, fixture.Store, fixture.AgentID, firstCompleted.ID) + firstResultEvents := listTypedToolResultEventsForToolCall( + t, + ctx, + fixture.Store, + fixture.AgentID, + firstCompleted.ID, + ) if len(firstResultEvents) != 1 { t.Fatalf("first tool result events = %d, want 1", len(firstResultEvents)) } @@ -515,7 +554,11 @@ func TestContinuationSeedKeepsLaterToolResultAfterEarlierResultConsumed(t *testi t, ctx, fixture, admitted.Turn.ID, laterContext.ID, "continuation_seed_later_result_second", fixture.Now.Add(8*time.Second), ) - if seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID); err != nil { + if seed, + found, + err := fixture.Store.Execution().NextAgentModelWork(ctx, + testProjectID, + fixture.AgentID); err != nil { t.Fatalf("next continuation seed: %v", err) } else if !found || seed.TurnID != completed.TurnID || len(seed.InputIDs) != 2 || seed.InputIDs[0] != admitted.Inputs[0].ID || seed.InputIDs[1] != admitted.Inputs[1].ID { @@ -526,56 +569,6 @@ func TestContinuationSeedKeepsLaterToolResultAfterEarlierResultConsumed(t *testi } } -func TestNextAgentContinuationSeedUsesOpeningInputsAtContextWatermark(t *testing.T) { - t.Parallel() - t.Run("orders multi-input opening events", func(t *testing.T) { - t.Parallel() - ctx := context.Background() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "continuation_seed_multi_input") - contextRecord := claimNormalContextAtFrontierTest( - t, ctx, fixture, []uuid.UUID{admitted.Inputs[0].ID, admitted.Inputs[1].ID}, agent.CurrentConfigID, - admitted.Events[1].Sequence, fixture.Now.Add(4*time.Second), - ) - completed := completeToolCallForContinuationSeedTest( - t, - ctx, - fixture, - admitted.Turn.ID, - contextRecord.ID, - "continuation_seed_multi_input", - fixture.Now.Add(5*time.Second), - ) - seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID) - if err != nil { - t.Fatalf("next continuation seed: %v", err) - } - if !found || seed.TurnID != completed.TurnID { - t.Fatalf("continuation seed found=%v seed=%+v want turn %s", found, seed, completed.TurnID) - } - if len(seed.InputIDs) != 2 || seed.InputIDs[0] != admitted.Inputs[0].ID || - seed.InputIDs[1] != admitted.Inputs[1].ID { - t.Fatalf("seed input order = %v, want [%s %s]", seed.InputIDs, admitted.Inputs[0].ID, admitted.Inputs[1].ID) - } - }) - - t.Run("rejects stale opening frontier", func(t *testing.T) { - t.Parallel() - ctx := context.Background() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "continuation_seed_watermark") - _, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{admitted.Inputs[0].ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) - if !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("stale opening frontier claim error = %v, want %v", err, storeerr.ErrAgentNotAdvanceable) - } - }) -} - func TestClaimNormalModelCallFencesContinuationToSelectedSource(t *testing.T) { t.Parallel() ctx := context.Background() @@ -635,37 +628,52 @@ func TestClaimNormalModelCallFencesContinuationToSelectedSource(t *testing.T) { if err != nil { t.Fatalf("load continuation frontier: %v", err) } - claim := func(sourceContextID, sourceOutputID uuid.UUID, now time.Time) (executionstore.ModelCallClaim, error) { - return fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: seed.InputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: frontier, - SourceModelCallContextID: sourceContextID, - SourceModelOutputID: sourceOutputID, - }) + claim := func(sourceContextID, + sourceOutputID uuid.UUID, + now time.Time) (executionstore.ModelCallClaim, + error) { + prepared, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: seed.InputIDs, + SourceModelCallContextID: sourceContextID, + SourceModelOutputID: sourceOutputID, + }) + return prepared.Claim, err } if _, err := claim(uuid.Nil, uuid.Nil, fixture.Now.Add(8*time.Second)); !errors.Is( err, storeerr.ErrAgentNotAdvanceable, ) { - t.Fatalf("continuation claim without source error = %v, want %v", err, storeerr.ErrAgentNotAdvanceable) + t.Fatalf( + "continuation claim without source error = %v, want %v", + err, + storeerr.ErrAgentNotAdvanceable, + ) } if _, err := claim( testID("wrong_continuation_source_context"), sourceOutput.ID, fixture.Now.Add(9*time.Second), ); !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("continuation claim with wrong source context error = %v, want %v", err, storeerr.ErrAgentNotAdvanceable) + t.Fatalf( + "continuation claim with wrong source context error = %v, want %v", + err, + storeerr.ErrAgentNotAdvanceable, + ) } if _, err := claim( sourceContext.ID, testID("wrong_continuation_source_output"), fixture.Now.Add(10*time.Second), ); !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("continuation claim with wrong source output error = %v, want %v", err, storeerr.ErrAgentNotAdvanceable) + t.Fatalf( + "continuation claim with wrong source output error = %v, want %v", + err, + storeerr.ErrAgentNotAdvanceable, + ) } created, err := claim(sourceContext.ID, sourceOutput.ID, fixture.Now.Add(11*time.Second)) @@ -739,13 +747,14 @@ func TestKernelContextEventsIncludesCanonicalTranscriptEvents(t *testing.T) { fixture.Now.Add(31*time.Second), ) - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"visible"}]`), - IdempotencyKey: "kernel-context-visible-input", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"visible"}]`), + IdempotencyKey: "kernel-context-visible-input", + }) if err != nil { t.Fatalf("create visible agent input: %v", err) } @@ -760,18 +769,15 @@ func TestKernelContextEventsIncludesCanonicalTranscriptEvents(t *testing.T) { if !found { t.Fatal("expected admitted visible agent input") } - agent, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) - if err != nil { - t.Fatalf("load agent for visible model output: %v", err) - } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, + }) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim visible model output context: %v", err) } @@ -868,7 +874,8 @@ func TestKernelContextEventsIncludesCanonicalTranscriptEvents(t *testing.T) { t.Fatalf("list compaction source events: %v", err) } for _, event := range compactionEvents { - if event.Sequence == 1 && strings.Contains(string(event.ContentParts), "Agent configuration changed") { + if event.Sequence == 1 && + strings.Contains(string(event.ContentParts), "Agent configuration changed") { t.Fatalf("baseline config change should not be projected into compaction source: %+v", event) } } @@ -919,13 +926,15 @@ func TestCompactionSourceEventsIncludeToolCallAndResultSemantics(t *testing.T) { siblingToolCall := toolCalls[1] markToolCallReadyForTest(t, ctx, fixture, toolCallID, now.Add(2*time.Millisecond)) markToolCallReadyForTest(t, ctx, fixture, siblingToolCall.ID, now.Add(3*time.Millisecond)) - openRows, err := fixture.Store.Execution().ListCompactionSourceEvents(ctx, testProjectID, fixture.AgentID, 0, 100) + openRows, err := fixture.Store.Execution(). + ListCompactionSourceEvents(ctx, testProjectID, fixture.AgentID, 0, 100) if err != nil { t.Fatalf("list open tool compaction source events: %v", err) } var openToolCallSequence int64 for _, row := range openRows { - if row.Kind == string(events.KindModelOutput) && strings.Contains(string(row.ContentParts), `"name": "run_command"`) { + if row.Kind == string(events.KindModelOutput) && + strings.Contains(string(row.ContentParts), `"name": "run_command"`) { openToolCallSequence = row.Sequence } } @@ -959,7 +968,8 @@ func TestCompactionSourceEventsIncludeToolCallAndResultSemantics(t *testing.T) { }); err != nil { t.Fatalf("complete tool call: %v", err) } - partialRows, err := fixture.Store.Execution().ListCompactionSourceEvents(ctx, testProjectID, fixture.AgentID, 0, 100) + partialRows, err := fixture.Store.Execution(). + ListCompactionSourceEvents(ctx, testProjectID, fixture.AgentID, 0, 100) if err != nil { t.Fatalf("list partially completed tool compaction source events: %v", err) } @@ -990,7 +1000,8 @@ func TestCompactionSourceEventsIncludeToolCallAndResultSemantics(t *testing.T) { }); err != nil { t.Fatalf("complete sibling tool call: %v", err) } - rows, err := fixture.Store.Execution().ListCompactionSourceEvents(ctx, testProjectID, fixture.AgentID, 0, 100) + rows, err := fixture.Store.Execution(). + ListCompactionSourceEvents(ctx, testProjectID, fixture.AgentID, 0, 100) if err != nil { t.Fatalf("list compaction source events: %v", err) } diff --git a/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go b/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go index 40b6f6ea6..d71dc7d5c 100644 --- a/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go @@ -17,9 +17,10 @@ import ( "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/toolcatalog" + "github.com/stretchr/testify/require" ) func listTypedToolResultEventsForToolCall( @@ -118,7 +119,7 @@ LIMIT 1 `, testProjectID, fixture.AgentID, turnID).Scan(&openingInputID, &inputSequence); err != nil { t.Fatalf("load opening event for model output: %v", err) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent for model output: %v", err) } @@ -163,16 +164,17 @@ LIMIT 1 if !found { t.Fatal("select model work for bare model output: no work found") } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{openingInputID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: inputSequence, - SourceModelCallContextID: work.ModelCallContextID, - SourceModelOutputID: work.SourceModelOutputID, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{openingInputID}, + SourceModelCallContextID: work.ModelCallContextID, + SourceModelOutputID: work.SourceModelOutputID, + }) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim bare model output context: %v", err) } @@ -398,16 +400,17 @@ func claimNormalContextAtFrontierTest( if !found { t.Fatal("select model work for context claim: no work found") } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agentConfigID, - InputEventSequence: inputEventSequence, - SourceModelCallContextID: work.ModelCallContextID, - SourceModelOutputID: work.SourceModelOutputID, - }) + prepared2, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, + SourceModelCallContextID: work.ModelCallContextID, + SourceModelOutputID: work.SourceModelOutputID, + }) + claim := prepared2.Claim + if err != nil { t.Fatalf("claim model context: %v", err) } @@ -481,6 +484,7 @@ func appendSyntheticLatestContentTurnForFrontierTest( if err := tx.Commit(ctx); err != nil { t.Fatalf("commit synthetic latest turn: %v", err) } + require.NoError(t, fixture.Store.Execution().RepairExecutionForTest(ctx, testProjectID, fixture.AgentID)) return turnID } @@ -513,14 +517,15 @@ func newStartedNormalModelCallTestFixture( for _, input := range admitted.Inputs { openingInputIDs = append(openingInputIDs, input.ID) } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[len(admitted.Events)-1].Sequence, - }) + prepared3, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, + }) + claim := prepared3.Claim + if err != nil { t.Fatalf("claim started normal model call: %v", err) } @@ -681,51 +686,51 @@ func claimSentCompactionForRangeTest( if err != nil { t.Fatalf("create checkpoint overflow input: %v", err) } - tx, err := fixture.Store.pool.Begin(ctx) - if err != nil { - t.Fatalf("begin checkpoint overflow admission: %v", err) - } - defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: testProjectID, - ID: fixture.AgentID, - }); err != nil { - t.Fatalf("lock checkpoint overflow agent: %v", err) - } - admitted, err := executionstore.IntegrationAdmitLockedAgentInputsAndOpenTurnTx( + unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) + require.NoError(t, err) + defer func() { _ = unit.Rollback(ctx) }() + _, err = unit.LockAgent( ctx, - notifications.NewTxNotifications(), - tx, - qtx, - executionstore.IntegrationAdmitAgentInputAndOpenTurnInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - }, - []executionstore.AgentInputRecord{overflowInput}, + dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: fixture.AgentID}, + agentexecution.ExternalAuthority{}, ) - if err != nil { - t.Fatalf("admit checkpoint overflow input: %v", err) + require.NoError(t, err) + admitted, err := executionstore.IntegrationAdmitInputs(ctx, unit, testProjectID, fixture.AgentID) + require.NoError(t, err) + if len(admitted.Events) == 0 { + require.NoError(t, unit.Rollback(ctx)) + id := uuid.New() + _, err = fixture.Store.pool.Exec( + ctx, + `INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,turn_id,operation_kind,attempt_number,agent_config_id,configured_model_revision_id,input_event_sequence,source_event_sequence_end,runtime_lock_id,opening_input_ids,opening_event_sequence,state,created_at) +SELECT $1,c.org_id,c.project_id,c.agent_id,c.turn_id,'compaction',1,c.agent_config_id, +c.configured_model_revision_id,$3,$3,$4,c.opening_input_ids,c.opening_event_sequence, +'started',statement_timestamp() +FROM agent_execution_state s JOIN model_call_contexts c ON c.agent_id=s.agent_id AND +c.id=s.normal_context_id WHERE s.agent_id=$2`, + id, + fixture.AgentID, + end, + fixture.Lock.ID, + ) + require.NoError(t, err) + record, found, err := fixture.Store.Execution(). + GetModelCallContext(ctx, testProjectID, fixture.AgentID, id) + require.NoError(t, err) + require.True(t, found) + return executionstore.ModelCallClaim{Context: record, Created: true, Claimed: true} } overflowFrontier := admitted.Events[len(admitted.Events)-1].Sequence - openingInputIDs, _, err := executionstore.IntegrationModelCallOpeningInputSet( - ctx, - qtx, - testProjectID, - fixture.AgentID, - admitted.Turn.ID, - overflowFrontier, - ) - if err != nil { - t.Fatalf("load checkpoint overflow opening inputs: %v", err) - } - if err := tx.Commit(ctx); err != nil { - t.Fatalf("commit checkpoint overflow admission: %v", err) - } + h, err := unit.Handle(testProjectID, fixture.AgentID) + require.NoError(t, err) + loaded, err := h.LoadExecution(ctx) + require.NoError(t, err) + openingInputIDs := loaded.Selection.Model.Opening.InputIDs + require.NoError(t, unit.Commit(ctx, "fixture compaction admission")) if len(admitted.Inputs) != 1 || admitted.Inputs[0].ID != overflowInput.ID { t.Fatalf("admitted checkpoint overflow input = %+v", admitted) } - snapshot, err := fixture.Store.Execution().CaptureAgentConfigForEventWatermark( + _, err = fixture.Store.Execution().CaptureAgentConfigForEventWatermark( ctx, testProjectID, fixture.AgentID, @@ -734,14 +739,15 @@ func claimSentCompactionForRangeTest( if err != nil { t.Fatalf("capture checkpoint overflow config: %v", err) } - parent, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: snapshot.AgentConfig.ID, - InputEventSequence: overflowFrontier, - }) + prepared4, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, + }) + parent := prepared4.Claim + if err != nil { t.Fatalf("claim checkpoint overflow parent: %v", err) } @@ -776,14 +782,15 @@ func publishCheckpointForRangeTest( summary string, now time.Time, ) (executionstore.ContextCheckpointRecord, error) { - return fixture.Store.Execution().PublishContextCheckpoint(ctx, executionstore.PublishContextCheckpointInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, ModelCallContextID: claim.Context.ID, - Summary: summary, - APIFormat: modelprotocol.APIFormatOpenAIResponses, - APIVariant: modelprotocol.APIVariantDefault, - ProviderResponseID: "resp_" + summary, - }) + return fixture.Store.Execution(). + PublishContextCheckpoint(ctx, executionstore.PublishContextCheckpointInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, ModelCallContextID: claim.Context.ID, + Summary: summary, + APIFormat: modelprotocol.APIFormatOpenAIResponses, + APIVariant: modelprotocol.APIVariantDefault, + ProviderResponseID: "resp_" + summary, + }) } func appendCancelStopEventForContinuationSeedTest( @@ -800,20 +807,16 @@ func appendCancelStopEventForContinuationSeedTest( t.Fatalf("begin cancel stop event fixture: %v", err) } defer func() { _ = tx.Rollback(ctx) }() - qtx := dbsqlc.New(tx) - controlType := "cancel_current" - controlInput, err := qtx.InsertControlAgentInput(ctx, dbsqlc.InsertControlAgentInputParams{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ActorID: storeutil.IDFromNil(actorID), - ControlType: &controlType, - IdempotencyScope: storeutil.TextFromEmpty("agent_control"), - InputIdempotencyKey: storeutil.TextFromEmpty("test-cancel-stop:" + turnID.String()), - Metadata: json.RawMessage(`{}`), - }) - if err != nil { - t.Fatalf("insert cancel control input: %v", err) - } + controlInput := executionstore.AgentInputRecord{ID: uuid.New()} + _, err = tx.Exec( + ctx, + `INSERT INTO agent_inputs(id,project_id,agent_id,actor_id,input_kind,delivery_mode,control_type,state,queued_at,metadata) VALUES($1,$2,$3,$4,'control','immediate','cancel_current','received',statement_timestamp(),'{}')`, + controlInput.ID, + testProjectID, + fixture.AgentID, + actorID, + ) + require.NoError(t, err) eventRecord, err := executionstore.IntegrationAppendTypedAgentEventTx( ctx, notifications.NewTxNotifications(), @@ -831,31 +834,29 @@ func appendCancelStopEventForContinuationSeedTest( t.Fatalf("append cancel stop event: %v", err) } if err := executionstore.IntegrationUpdateAgentTurnLatestEventQuery( - ctx, qtx, testProjectID, fixture.AgentID, turnID, eventRecord.Event.ID, uuid.Nil, + ctx, tx, testProjectID, fixture.AgentID, turnID, eventRecord.Event.ID, uuid.Nil, ); err != nil { t.Fatalf("update turn latest cancel event: %v", err) } - if err := qtx.ResolveControlAgentInput( + _, err = tx.Exec( ctx, - dbsqlc.ResolveControlAgentInputParams{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ID: controlInput.ID, - ControlType: &controlType, - EventID: &eventRecord.Event.ID, - }, - ); err != nil { - t.Fatalf("resolve cancel control input: %v", err) - } + `UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=$3,resolved_at=$3 WHERE id=$1`, + controlInput.ID, + eventRecord.Event.ID, + eventRecord.Event.At, + ) + require.NoError(t, err) if err := tx.Commit(ctx); err != nil { t.Fatalf("commit cancel stop event fixture: %v", err) } + require.NoError(t, fixture.Store.Execution().RepairExecutionForTest(ctx, testProjectID, fixture.AgentID)) return eventRecord.Event } func isPgConstraintViolation(err error) bool { var pgErr *pgconn.PgError - return errors.As(err, &pgErr) && (pgErr.Code == "23505" || pgErr.Code == "23514" || pgErr.Code == "23503") + return errors.As(err, &pgErr) && + (pgErr.Code == "23502" || pgErr.Code == "23505" || pgErr.Code == "23514" || pgErr.Code == "23503") } func assertPgErrorMessage(t *testing.T, err error, code, message string) { diff --git a/internal/storage/executionstore/kernel_event_authority_model_calls_integration_test.go b/internal/storage/executionstore/kernel_event_authority_model_calls_integration_test.go index deaf527ea..c75b8de55 100644 --- a/internal/storage/executionstore/kernel_event_authority_model_calls_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_model_calls_integration_test.go @@ -20,25 +20,25 @@ import ( func TestModelCallRowConstraintsProtectImmutableEvidence(t *testing.T) { t.Parallel() ctx := context.Background() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "model_call_database_guards") + fixture, admitted, _ := newMultiInputContinuationSeedFixture(t, ctx, "model_call_database_guards") openingInputIDs := make([]uuid.UUID, 0, len(admitted.Inputs)) for _, input := range admitted.Inputs { openingInputIDs = append(openingInputIDs, input.ID) } - frontier := admitted.Events[len(admitted.Events)-1].Sequence - claimInput := executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: frontier, + _ = admitted.Events[len(admitted.Events)-1].Sequence + claimInput := executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, claimInput) + prepared1, err := fixture.Store.Execution().PrepareNormalModelCall(ctx, claimInput) + claim := prepared1.Claim if err != nil { t.Fatalf("claim guarded model context: %v", err) } - replayed, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, claimInput) + prepared2, err := fixture.Store.Execution().PrepareNormalModelCall(ctx, claimInput) + replayed := prepared2.Claim if err != nil { t.Fatalf("re-find normal model context: %v", err) } @@ -47,12 +47,13 @@ func TestModelCallRowConstraintsProtectImmutableEvidence(t *testing.T) { } _, terminalInsertErr := fixture.Store.pool.Exec(ctx, ` INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, + turn_id, opening_input_ids, opening_event_sequence, +org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, runtime_lock_id, state, error_kind, error_message, created_at, completed_at ) -SELECT org_id, project_id, agent_id, operation_kind, +SELECT turn_id, opening_input_ids, opening_event_sequence, org_id, project_id, agent_id, operation_kind, attempt_number + 1000, agent_config_id, configured_model_revision_id, input_event_sequence, runtime_lock_id, 'failed', 'invalid_request', 'inserted terminal', created_at, created_at @@ -71,13 +72,14 @@ WHERE id = $1`, claim.Context.ID) } _, nullRangeErr := nullRangeTx.Exec(ctx, ` INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, + turn_id, opening_input_ids, opening_event_sequence, +org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, source_event_sequence_end, runtime_lock_id, state, created_at ) -SELECT org_id, project_id, agent_id, 'compaction', +SELECT turn_id, opening_input_ids, opening_event_sequence, org_id, project_id, agent_id, 'compaction', attempt_number + 1001, agent_config_id, configured_model_revision_id, input_event_sequence, NULL, runtime_lock_id, 'started', created_at @@ -166,7 +168,11 @@ func TestKernelRecordModelOutputWritesTypedAuthority(t *testing.T) { t.Run(string(stopReason), func(t *testing.T) { t.Parallel() ctx := t.Context() - fixture, _, modelClaim := newStartedNormalModelCallTestFixture(t, ctx, "kernel_model_output_authority") + fixture, _, modelClaim := newStartedNormalModelCallTestFixture( + t, + ctx, + "kernel_model_output_authority", + ) now := fixture.Now.Add(time.Minute) providerModelSlug := modelProviderSlugForContext( t, ctx, fixture.Store, testProjectID, fixture.AgentID, modelClaim.Context.ID, @@ -214,12 +220,19 @@ func TestKernelRecordModelOutputWritesTypedAuthority(t *testing.T) { t.Fatalf("replay typed model output: %v", err) } if replayed.ID != event.ID || replayed.Sequence != event.Sequence { - t.Fatalf("replayed event = %s/%d, want %s/%d", replayed.ID, replayed.Sequence, event.ID, event.Sequence) + t.Fatalf( + "replayed event = %s/%d, want %s/%d", + replayed.ID, + replayed.Sequence, + event.ID, + event.Sequence, + ) } changedReason := recordInput changedReason.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonToolUse if _, err := fixture.Store.Execution().RecordModelOutputAndCompleteContext(ctx, changedReason); !errors.Is( - err, storeerr.ErrIdempotencyConflict, + err, + storeerr.ErrIdempotencyConflict, ) { t.Fatalf("changed stop reason replay error = %v, want idempotency conflict", err) } @@ -227,7 +240,10 @@ func TestKernelRecordModelOutputWritesTypedAuthority(t *testing.T) { require.NoError( t, fixture.Store.pool.QueryRow( - ctx, `SELECT count(*) FROM agent_next_model_work($1,$2)`, testProjectID, fixture.AgentID, + ctx, + `SELECT count(*) FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id WHERE a.project_id=$1 AND h.agent_id=$2 AND h.logical_ready_at IS NOT NULL`, + testProjectID, + fixture.AgentID, ). Scan(&pending), ) @@ -249,7 +265,11 @@ func TestKernelRecordModelOutputWritesTypedAuthority(t *testing.T) { err, storeerr.ErrIdempotencyConflict, ) { - t.Fatalf("conflicting output replay error = %v, want %v", err, storeerr.ErrIdempotencyConflict) + t.Fatalf( + "conflicting output replay error = %v, want %v", + err, + storeerr.ErrIdempotencyConflict, + ) } conflictingCostInput := recordInput conflictingCostInput.ProviderResponse.ProviderReportedCostUSD = "0.0000126" @@ -263,7 +283,10 @@ func TestKernelRecordModelOutputWritesTypedAuthority(t *testing.T) { var modelOutputID uuid.UUID if err := fixture.Store.pool.QueryRow(ctx, `SELECT event.model_output_id FROM agent_events event JOIN agents agent ON agent.id = event.agent_id -WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.id = $3`, testProjectID, fixture.AgentID, event.ID). +WHERE agent.project_id = $1 AND event.agent_id = $2 AND event.id = $3`, + testProjectID, + fixture.AgentID, + event.ID). Scan(&modelOutputID); err != nil { t.Fatalf("load model output event pointer: %v", err) } @@ -348,7 +371,12 @@ JOIN model_call_contexts context ON context.agent_id = output.agent_id AND context.id = output.model_call_context_id WHERE context.project_id = $1 AND output.agent_id = $2 AND output.id = $3 - `, testProjectID, fixture.AgentID, modelOutputID).Scan(&inputTokens, &outputTokens, &contentBlocks); err != nil { + `, + testProjectID, + fixture.AgentID, + modelOutputID).Scan(&inputTokens, + &outputTokens, + &contentBlocks); err != nil { t.Fatalf("load model output usage: %v", err) } if inputTokens != 1 || outputTokens != 1 || contentBlocks != 1 { diff --git a/internal/storage/executionstore/kernel_event_authority_model_outputs.go b/internal/storage/executionstore/kernel_event_authority_model_outputs.go index e4c05ea50..72739233e 100644 --- a/internal/storage/executionstore/kernel_event_authority_model_outputs.go +++ b/internal/storage/executionstore/kernel_event_authority_model_outputs.go @@ -1,7 +1,6 @@ package executionstore import ( - "bytes" "context" "encoding/json" "errors" @@ -12,7 +11,6 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/storeerr" ) type CreateModelOutputAuthorityInput struct { @@ -44,7 +42,9 @@ func (s *Store) GetModelOutputForContext( projectID, agentID, modelCallContextID uuid.UUID, ) (ModelOutputAuthorityRecord, bool, error) { if projectID == uuid.Nil || agentID == uuid.Nil || modelCallContextID == uuid.Nil { - return ModelOutputAuthorityRecord{}, false, errors.New("project, agent, and model context are required") + return ModelOutputAuthorityRecord{}, false, errors.New( + "project, agent, and model context are required", + ) } row, err := s.q.GetModelOutputByModelContext(ctx, dbsqlc.GetModelOutputByModelContextParams{ ProjectID: projectID, @@ -59,57 +59,3 @@ func (s *Store) GetModelOutputForContext( } return modelOutputAuthorityFromGetSQLC(row), true, nil } - -func createModelOutputAuthorityTx( - ctx context.Context, - db sqlExecutor, - input CreateModelOutputAuthorityInput, -) (ModelOutputAuthorityRecord, error) { - if err := validateModelOutputAuthorityInput(input); err != nil { - return ModelOutputAuthorityRecord{}, err - } - input.Usage = modelUsageForStorage(input.Usage) - providerReplay := bytes.TrimSpace(input.ProviderReplay) - if len(providerReplay) == 0 || bytes.Equal(providerReplay, []byte("null")) { - providerReplay = nil - } else if !json.Valid(providerReplay) { - return ModelOutputAuthorityRecord{}, errors.New("provider replay must be valid JSON") - } - input.ProviderReplay = providerReplay - q := dbsqlc.New(db) - row, err := q.InsertModelOutputAuthority(ctx, dbsqlc.InsertModelOutputAuthorityParams{ - ModelCallContextID: input.ModelCallContextID, - ServedProviderModelSlug: input.ServedProviderModelSlug, - StopReason: string(input.StopReason), - ProviderReplay: sqlcRawMessageFromEmpty(input.ProviderReplay), - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }) - if errors.Is(err, pgx.ErrNoRows) { - existing, getErr := q.GetModelOutputByModelContext( - ctx, - dbsqlc.GetModelOutputByModelContextParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - }, - ) - if getErr != nil { - return ModelOutputAuthorityRecord{}, fmt.Errorf( - "load model output authority: %w", - getErr, - ) - } - record := modelOutputAuthorityFromGetSQLC(existing) - if !sameModelOutputAuthorityIntent(record, input) { - return ModelOutputAuthorityRecord{}, storeerr.ErrIdempotencyConflict - } - return record, nil - } - if err != nil { - return ModelOutputAuthorityRecord{}, fmt.Errorf("create model output authority: %w", err) - } - record := modelOutputAuthorityFromSQLC(row) - record.Usage = input.Usage - return record, nil -} diff --git a/internal/storage/executionstore/kernel_event_authority_store.go b/internal/storage/executionstore/kernel_event_authority_store.go index 0d5a83918..b41161738 100644 --- a/internal/storage/executionstore/kernel_event_authority_store.go +++ b/internal/storage/executionstore/kernel_event_authority_store.go @@ -1,22 +1,14 @@ package executionstore import ( - "context" "encoding/json" - "errors" - "fmt" - "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/dbsafe" "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/resourcemeta" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" ) type TypedAgentEventRecord struct { @@ -46,550 +38,6 @@ type CreateContentBlockInput struct { Metadata resourcemeta.Metadata } -type ContentBlockRecord struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - OwnerKind ContentBlockOwnerKind - OwnerAgentInputID uuid.UUID - OwnerModelOutputID uuid.UUID - OwnerToolCallResultID uuid.UUID - Ordinal int32 - BlockKind ContentBlockKind - TextContent string - StructuredData json.RawMessage - ArtifactID uuid.UUID - ToolCallID uuid.UUID - CreatedAt time.Time -} - -type admittedToolCallResult struct { - Result ToolCallResultAuthorityRecord - Event TypedAgentEventRecord - ContentBlocks json.RawMessage - Inserted bool -} - -type AppendTypedAgentEventInput struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - IsOpeningEvent bool - Kind events.Kind - IdempotencyKey string - AgentInputID uuid.UUID - ModelOutputID uuid.UUID - ToolCallResultID uuid.UUID - ContextCheckpointID uuid.UUID -} - -type sqlExecutor interface { - dbsqlc.DBTX -} - -func createContentBlockTx( - ctx context.Context, - db sqlExecutor, - input CreateContentBlockInput, -) (ContentBlockRecord, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || - input.OwnerKind == "" || - input.BlockKind == "" { - return ContentBlockRecord{}, errors.New( - "project, agent, owner kind, and block kind are required", - ) - } - if err := dbsafe.Text(input.TextContent); err != nil { - return ContentBlockRecord{}, fmt.Errorf("content block text %w", err) - } - if len(input.StructuredData) > 0 { - normalized, err := normalizeContentBlockJSONForStorage(input.StructuredData) - if err != nil { - return ContentBlockRecord{}, fmt.Errorf("content block structured data %w", err) - } - input.StructuredData = normalized - } - metadata, err := input.Metadata.JSON() - if err != nil { - return ContentBlockRecord{}, err - } - var textContent *string - if input.BlockKind == ContentBlockKindText || - input.BlockKind == ContentBlockKindReasoning || - input.BlockKind == ContentBlockKindError || - input.TextContent != "" { - textContent = &input.TextContent - } - row, err := dbsqlc.New(db).InsertContentBlock(ctx, dbsqlc.InsertContentBlockParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - OwnerKind: string(input.OwnerKind), - OwnerAgentInputID: storeutil.IDFromNil(input.OwnerAgentInputID), - OwnerModelOutputID: storeutil.IDFromNil(input.OwnerModelOutputID), - OwnerToolCallResultID: storeutil.IDFromNil(input.OwnerToolCallResultID), - Ordinal: input.Ordinal, - BlockKind: string(input.BlockKind), - TextContent: textContent, - StructuredData: sqlcRawMessageFromEmpty(input.StructuredData), - ArtifactID: storeutil.IDFromNil(input.ArtifactID), - ToolCallID: storeutil.IDFromNil(input.ToolCallID), - ExcludeFromModelContext: input.ExcludeFromModelContext, - Metadata: metadata, - }) - if err != nil { - return ContentBlockRecord{}, fmt.Errorf("create content block: %w", err) - } - return contentBlockFromInsertSQLC(row), nil -} - -func appendTypedAgentEventTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - input AppendTypedAgentEventInput, -) (TypedAgentEventRecord, error) { - if input.TurnID == uuid.Nil { - return TypedAgentEventRecord{}, errors.New("turn id is required") - } - if err := validateTypedEventPointers(input); err != nil { - return TypedAgentEventRecord{}, err - } - qtx := dbsqlc.New(tx) - allocation, err := qtx.AllocateEventSequence(ctx, dbsqlc.AllocateEventSequenceParams{ID: input.AgentID}) - if err != nil { - return TypedAgentEventRecord{}, fmt.Errorf("allocate event sequence: %w", err) - } - if input.ProjectID != allocation.ProjectID { - return TypedAgentEventRecord{}, storeerr.ErrNotFound - } - if input.IdempotencyKey != "" { - existing, found, err := loadTypedEventByIdempotencyMaybeTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.IdempotencyKey, - ) - if err != nil { - return TypedAgentEventRecord{}, err - } - if found { - if !sameTypedEventIntent(existing, input) { - return TypedAgentEventRecord{}, storeerr.ErrIdempotencyConflict - } - return existing, err - } - } - row, err := qtx.InsertTypedAgentEvent(ctx, dbsqlc.InsertTypedAgentEventParams{ - ID: storeutil.IDFromNil(input.ID), - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: input.TurnID, - Sequence: allocation.NextEventSequence, - EventKind: string(input.Kind), - IdempotencyKey: storeutil.TextFromEmpty(input.IdempotencyKey), - AgentInputID: storeutil.IDFromNil(input.AgentInputID), - ModelOutputID: storeutil.IDFromNil(input.ModelOutputID), - ToolCallResultID: storeutil.IDFromNil(input.ToolCallResultID), - ContextCheckpointID: storeutil.IDFromNil(input.ContextCheckpointID), - IsOpeningEvent: input.IsOpeningEvent, - }) - if errors.Is(err, pgx.ErrNoRows) { - return TypedAgentEventRecord{}, storeerr.ErrIdempotencyConflict - } - if err != nil { - return TypedAgentEventRecord{}, fmt.Errorf("insert typed agent event: %w", err) - } - if err := qtx.AdvanceEventSequence( - ctx, - dbsqlc.AdvanceEventSequenceParams{ProjectID: input.ProjectID, ID: input.AgentID}, - ); err != nil { - return TypedAgentEventRecord{}, fmt.Errorf("advance event sequence: %w", err) - } - txNotifications.AddAgentEvent(input.AgentID, allocation.NextEventSequence, string(input.Kind)) - return typedAgentEventFromInsertSQLC(row) -} - -func loadTypedEventByIdempotencyMaybeTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID uuid.UUID, - key string, -) (TypedAgentEventRecord, bool, error) { - row, err := dbsqlc.New(tx). - GetTypedAgentEventByIdempotency(ctx, dbsqlc.GetTypedAgentEventByIdempotencyParams{ - ProjectID: projectID, - AgentID: agentID, - IdempotencyKey: key, - }) - if errors.Is(err, pgx.ErrNoRows) { - return TypedAgentEventRecord{}, false, nil - } - if err != nil { - return TypedAgentEventRecord{}, false, fmt.Errorf( - "load typed event by idempotency: %w", - err, - ) - } - record, err := typedAgentEventFromIdempotencySQLC(row) - return record, true, err -} - -func loadTypedEventByModelOutputMaybeTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID, modelOutputID uuid.UUID, -) (TypedAgentEventRecord, bool, error) { - row, err := dbsqlc.New(tx).GetTypedAgentEventByModelOutput( - ctx, - dbsqlc.GetTypedAgentEventByModelOutputParams{ - ProjectID: projectID, - AgentID: agentID, - ModelOutputID: storeutil.IDFromNil(modelOutputID), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return TypedAgentEventRecord{}, false, nil - } - if err != nil { - return TypedAgentEventRecord{}, false, fmt.Errorf("load typed event by model output: %w", err) - } - record, err := typedAgentEventFromModelOutputSQLC(row) - return record, true, err -} - -type CreateToolCallResultAuthorityInput struct { - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - ToolCallID uuid.UUID - Outcome ToolResultOutcome - ResultContentParts json.RawMessage - IdempotencyKey string -} - -type ToolCallResultAuthorityRecord struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - ToolCallID uuid.UUID - Outcome ToolResultOutcome - CompletedAt time.Time -} - -func createToolCallResultAuthorityTx( - ctx context.Context, - db sqlExecutor, - input CreateToolCallResultAuthorityInput, -) (ToolCallResultAuthorityRecord, bool, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ToolCallID == uuid.Nil { - return ToolCallResultAuthorityRecord{}, false, errors.New( - "project, agent, and tool call are required", - ) - } - if !input.Outcome.IsTerminal() { - return ToolCallResultAuthorityRecord{}, false, errors.New( - "terminal tool result outcome is required", - ) - } - row, err := dbsqlc.New(db). - InsertToolCallResultAuthority(ctx, dbsqlc.InsertToolCallResultAuthorityParams{ - Outcome: string(input.Outcome), - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ToolCallID: input.ToolCallID, - }) - if errors.Is(err, pgx.ErrNoRows) { - existing, found, getErr := getToolCallResultAuthorityByToolCallTx( - ctx, - db, - input.ProjectID, - input.AgentID, - input.ToolCallID, - ) - if getErr != nil { - return ToolCallResultAuthorityRecord{}, false, getErr - } - if !found { - return ToolCallResultAuthorityRecord{}, false, storeerr.ErrStateTransitionConflict - } - if !sameToolCallResultIntent(existing, input) { - return ToolCallResultAuthorityRecord{}, false, storeerr.ErrIdempotencyConflict - } - existingParts, partsErr := toolCallResultContentBlocksTx( - ctx, - dbsqlc.New(db), - input.ProjectID, - input.AgentID, - existing.ID, - ) - if partsErr != nil { - return ToolCallResultAuthorityRecord{}, false, partsErr - } - if !sameJSON(existingParts, input.ResultContentParts) { - return ToolCallResultAuthorityRecord{}, false, storeerr.ErrIdempotencyConflict - } - return existing, false, nil - } - if err != nil { - return ToolCallResultAuthorityRecord{}, false, fmt.Errorf( - "create tool call result authority: %w", - err, - ) - } - return toolCallResultAuthorityFromInsertSQLC(row), true, nil -} - -func admitToolCallResultTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - input CreateToolCallResultAuthorityInput, -) (admittedToolCallResult, error) { - if len(input.ResultContentParts) == 0 { - input.ResultContentParts = json.RawMessage(`[]`) - } - blocks, err := parseToolResultContentBlocks(input.ResultContentParts) - if err != nil { - return admittedToolCallResult{}, err - } - canonical, err := marshalToolResultContentBlocks(blocks) - if err != nil { - return admittedToolCallResult{}, err - } - input.ResultContentParts = canonical - result, inserted, err := createToolCallResultAuthorityTx(ctx, tx, input) - if err != nil { - return admittedToolCallResult{}, err - } - if inserted { - for _, block := range blocks { - block.ProjectID = input.ProjectID - block.AgentID = input.AgentID - block.OwnerKind = ContentBlockOwnerToolCallResult - block.OwnerToolCallResultID = result.ID - if _, err := createContentBlockTx(ctx, tx, block); err != nil { - return admittedToolCallResult{}, err - } - } - } - event, err := appendTypedAgentEventTx( - ctx, - txNotifications, - tx, - AppendTypedAgentEventInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: input.TurnID, - Kind: events.KindToolResult, - IdempotencyKey: input.IdempotencyKey, - ToolCallResultID: result.ID, - }, - ) - if err != nil { - return admittedToolCallResult{}, err - } - return admittedToolCallResult{ - Result: result, - Event: event, - ContentBlocks: canonical, - Inserted: inserted, - }, nil -} - -func toolCallResultContentBlocksTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, resultID uuid.UUID, -) (json.RawMessage, error) { - rows, err := qtx.ListToolCallResultContentBlocks( - ctx, - dbsqlc.ListToolCallResultContentBlocksParams{ - ProjectID: projectID, - AgentID: agentID, - ToolCallResultID: storeutil.IDFromNil(resultID), - }, - ) - if err != nil { - return nil, fmt.Errorf("list tool result content blocks: %w", err) - } - parts := make([]map[string]any, 0, len(rows)) - for _, row := range rows { - decoded, err := resourcemeta.FromJSON(row.Metadata) - if err != nil { - return nil, fmt.Errorf("stored tool result content block metadata: %w", err) - } - metadata := contentBlockMetadataForOutput(decoded) - var part map[string]any - switch ContentBlockKind(row.BlockKind) { - case ContentBlockKindText: - part = map[string]any{"type": "text", "text": row.TextContent} - case ContentBlockKindStructuredData: - if row.StructuredData == nil { - return nil, errors.New("stored structured_data block has no value") - } - part = map[string]any{ - "type": "structured_data", - "value": *row.StructuredData, - } - case ContentBlockKindArtifact: - if row.ArtifactID == nil { - return nil, errors.New("stored artifact block has no artifact") - } - part = map[string]any{ - "type": "media_ref", - "artifact_id": row.ArtifactID.String(), - } - if row.ExcludeFromModelContext { - part["exclude_from_model_context"] = true - } - default: - return nil, fmt.Errorf("stored tool result has unsupported block kind %q", row.BlockKind) - } - if metadata != nil { - part["metadata"] = metadata - } - parts = append(parts, part) - } - body, err := marshalJSON(parts) - if err != nil { - return nil, fmt.Errorf("marshal stored tool result content parts: %w", err) - } - return normalizedJSON(body), nil -} - -func getToolCallResultAuthorityByToolCallTx( - ctx context.Context, - db sqlExecutor, - projectID, agentID, toolCallID uuid.UUID, -) (ToolCallResultAuthorityRecord, bool, error) { - row, err := dbsqlc.New(db). - GetToolCallResultByToolCall(ctx, dbsqlc.GetToolCallResultByToolCallParams{ - ProjectID: projectID, - AgentID: agentID, - ToolCallID: toolCallID, - }) - if errors.Is(err, pgx.ErrNoRows) { - return ToolCallResultAuthorityRecord{}, false, nil - } - if err != nil { - return ToolCallResultAuthorityRecord{}, false, fmt.Errorf( - "get tool call result authority: %w", - err, - ) - } - return toolCallResultAuthorityFromGetSQLC(row), true, nil -} - -func validateModelOutputAuthorityInput(input CreateModelOutputAuthorityInput) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ModelCallContextID == uuid.Nil || - input.StopReason == "" { - return errors.New( - "project, agent, model context, and stop reason are required", - ) - } - if !modelenvelope.IsDurableModelOutputStopReason(input.StopReason) { - return fmt.Errorf("unsupported model output stop reason %q", input.StopReason) - } - return nil -} - -func sameTypedEventIntent(existing TypedAgentEventRecord, input AppendTypedAgentEventInput) bool { - return existing.Event.Kind == input.Kind && - existing.Event.IdempotencyKey == input.IdempotencyKey && - existing.TurnID == input.TurnID && - existing.IsOpeningEvent == input.IsOpeningEvent && - existing.AgentInputID == input.AgentInputID && - existing.ModelOutputID == input.ModelOutputID && - existing.ToolCallResultID == input.ToolCallResultID && - existing.ContextCheckpointID == input.ContextCheckpointID -} - -func sameToolCallResultIntent( - existing ToolCallResultAuthorityRecord, - input CreateToolCallResultAuthorityInput, -) bool { - return existing.ToolCallID == input.ToolCallID && - existing.Outcome == input.Outcome -} - -func validateTypedEventPointers(input AppendTypedAgentEventInput) error { - count := 0 - if input.AgentInputID != uuid.Nil { - count++ - } - if input.ModelOutputID != uuid.Nil { - count++ - } - if input.ToolCallResultID != uuid.Nil { - count++ - } - if input.ContextCheckpointID != uuid.Nil { - count++ - } - if count != 1 { - return errors.New("typed agent event requires exactly one typed pointer") - } - switch input.Kind { - case events.KindAgentInput: - if input.AgentInputID == uuid.Nil { - return errors.New("agent_input event requires agent_input_id") - } - case events.KindModelOutput: - if input.ModelOutputID == uuid.Nil { - return errors.New("model_output event requires model_output_id") - } - case events.KindToolResult: - if input.ToolCallResultID == uuid.Nil { - return errors.New("tool_result event requires tool_call_result_id") - } - case events.KindContextCheckpoint: - if input.ContextCheckpointID == uuid.Nil { - return errors.New("context_checkpoint event requires context_checkpoint_id") - } - default: - return fmt.Errorf("invalid typed frontier event kind: %s", input.Kind) - } - return nil -} - -func sameModelOutputAuthorityIntent( - existing ModelOutputAuthorityRecord, - input CreateModelOutputAuthorityInput, -) bool { - return existing.ModelCallContextID == input.ModelCallContextID && - existing.ServedProviderModelSlug == input.ServedProviderModelSlug && - existing.StopReason == input.StopReason && - ((len(existing.ProviderReplay) == 0 && len(input.ProviderReplay) == 0) || - sameJSON(existing.ProviderReplay, input.ProviderReplay)) && - existing.Usage == modelUsageForStorage(input.Usage) -} - -func modelOutputAuthorityFromSQLC(row dbsqlc.InsertModelOutputAuthorityRow) ModelOutputAuthorityRecord { - return ModelOutputAuthorityRecord{ - ID: row.ID, - ProjectID: row.ProjectID, - AgentID: row.AgentID, - TurnID: row.TurnID, - ModelCallContextID: row.ModelCallContextID, - ServedProviderModelSlug: row.ServedProviderModelSlug, - StopReason: modelenvelope.StopReason(row.StopReason), - ProviderResponseID: row.ProviderResponseID, - ProviderReplay: rawMessageFromSQLCPtr(row.ProviderReplay), - Usage: modelUsageFromSQLC( - row.InputTokensTotal, - row.UncachedInputTokens, - row.CacheReadInputTokens, - row.CacheWriteInputTokens, - row.OutputTokensTotal, - row.ReasoningOutputTokens, - ), - CreatedAt: row.CreatedAt, - } -} - func modelOutputAuthorityFromGetSQLC(row dbsqlc.GetModelOutputByModelContextRow) ModelOutputAuthorityRecord { return ModelOutputAuthorityRecord{ ID: row.ID, @@ -613,103 +61,6 @@ func modelOutputAuthorityFromGetSQLC(row dbsqlc.GetModelOutputByModelContextRow) } } -func toolCallResultAuthorityFromInsertSQLC(row dbsqlc.InsertToolCallResultAuthorityRow) ToolCallResultAuthorityRecord { - return ToolCallResultAuthorityRecord{ - ID: row.ID, - ProjectID: row.ProjectID, - AgentID: row.AgentID, - TurnID: row.TurnID, - ToolCallID: row.ToolCallID, - Outcome: ToolResultOutcome(row.Outcome), - CompletedAt: row.CompletedAt, - } -} - -func toolCallResultAuthorityFromGetSQLC(row dbsqlc.GetToolCallResultByToolCallRow) ToolCallResultAuthorityRecord { - return ToolCallResultAuthorityRecord{ - ID: row.ID, - ProjectID: row.ProjectID, - AgentID: row.AgentID, - TurnID: row.TurnID, - ToolCallID: row.ToolCallID, - Outcome: ToolResultOutcome(row.Outcome), - CompletedAt: row.CompletedAt, - } -} - -func contentBlockFromInsertSQLC(row dbsqlc.InsertContentBlockRow) ContentBlockRecord { - return ContentBlockRecord{ - ID: row.ID, - ProjectID: row.ProjectID, - AgentID: row.AgentID, - OwnerKind: ContentBlockOwnerKind(row.OwnerKind), - OwnerAgentInputID: storeutil.IDFromPtr(row.OwnerAgentInputID), - OwnerModelOutputID: storeutil.IDFromPtr(row.OwnerModelOutputID), - OwnerToolCallResultID: storeutil.IDFromPtr(row.OwnerToolCallResultID), - Ordinal: row.Ordinal, - BlockKind: ContentBlockKind(row.BlockKind), - TextContent: row.TextContent, - StructuredData: rawMessageFromSQLCPtr(row.StructuredData), - ArtifactID: storeutil.IDFromPtr(row.ArtifactID), - ToolCallID: storeutil.IDFromPtr(row.ToolCallID), - CreatedAt: row.CreatedAt, - } -} - -func typedAgentEventFromInsertSQLC( - row dbsqlc.InsertTypedAgentEventRow, -) (TypedAgentEventRecord, error) { - event, err := events.New( - events.NewInput{ - ID: row.ID, - AgentID: row.AgentID, - Sequence: row.Sequence, - Kind: events.Kind(row.EventKind), - At: row.CreatedAt, - IdempotencyKey: row.IdempotencyKey, - }, - ) - if err != nil { - return TypedAgentEventRecord{}, err - } - return TypedAgentEventRecord{ - Event: event, - TurnID: row.TurnID, - IsOpeningEvent: row.IsOpeningEvent, - AgentInputID: storeutil.IDFromPtr(row.AgentInputID), - ModelOutputID: storeutil.IDFromPtr(row.ModelOutputID), - ToolCallResultID: storeutil.IDFromPtr(row.ToolCallResultID), - ContextCheckpointID: storeutil.IDFromPtr(row.ContextCheckpointID), - }, nil -} - -func typedAgentEventFromIdempotencySQLC( - row dbsqlc.GetTypedAgentEventByIdempotencyRow, -) (TypedAgentEventRecord, error) { - event, err := events.New( - events.NewInput{ - ID: row.ID, - AgentID: row.AgentID, - Sequence: row.Sequence, - Kind: events.Kind(row.EventKind), - At: row.CreatedAt, - IdempotencyKey: row.IdempotencyKey, - }, - ) - if err != nil { - return TypedAgentEventRecord{}, err - } - return TypedAgentEventRecord{ - Event: event, - TurnID: row.TurnID, - IsOpeningEvent: row.IsOpeningEvent, - AgentInputID: storeutil.IDFromPtr(row.AgentInputID), - ModelOutputID: storeutil.IDFromPtr(row.ModelOutputID), - ToolCallResultID: storeutil.IDFromPtr(row.ToolCallResultID), - ContextCheckpointID: storeutil.IDFromPtr(row.ContextCheckpointID), - }, nil -} - func typedAgentEventFromModelOutputSQLC( row dbsqlc.GetTypedAgentEventByModelOutputRow, ) (TypedAgentEventRecord, error) { diff --git a/internal/storage/executionstore/machine_daemon_bootstrap_store.go b/internal/storage/executionstore/machine_daemon_bootstrap_store.go index 47f6a10b0..8a7a5c434 100644 --- a/internal/storage/executionstore/machine_daemon_bootstrap_store.go +++ b/internal/storage/executionstore/machine_daemon_bootstrap_store.go @@ -15,6 +15,7 @@ import ( "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/resourcename" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -332,30 +333,31 @@ func (s *Store) RevokeBYOMachineDaemonToken( orgID, machineID, tokenID uuid.UUID, reason string, ) (MachineDaemonTokenRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return MachineDaemonTokenRecord{}, fmt.Errorf("begin revoke machine daemon token: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() - record, err := s.RevokeBYOMachineDaemonTokenTx(ctx, tx, txNotifications, orgID, machineID, tokenID, reason) + defer func() { _ = unit.Rollback(ctx) }() + + record, err := s.RevokeBYOMachineDaemonTokenInUnit(ctx, unit, orgID, machineID, tokenID, reason) if err != nil { return MachineDaemonTokenRecord{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "revoke machine daemon token"); err != nil { + if err := unit.Commit(ctx, "revoke machine daemon token"); err != nil { return MachineDaemonTokenRecord{}, err } return record, nil } -func (s *Store) RevokeBYOMachineDaemonTokenTx( +func (s *Store) RevokeBYOMachineDaemonTokenInUnit( ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, + unit *agentexecution.Unit, orgID, machineID, tokenID uuid.UUID, reason string, ) (MachineDaemonTokenRecord, error) { - qtx := s.q.WithTx(tx) + txNotifications := unit.Notifications() + tx := unit.DB() + qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForRuntimeRegistration( ctx, dbsqlc.LockMachineForRuntimeRegistrationParams{OrgID: orgID, ID: machineID}, @@ -411,8 +413,7 @@ func (s *Store) RevokeBYOMachineDaemonTokenTx( ) if err := completeMachineLifecycleTerminalWorkTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, machineID, @@ -543,13 +544,14 @@ func (s *Store) RecordMachineFailureReport( MachineID: input.MachineID, DaemonTokenID: input.DaemonTokenID, } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return fmt.Errorf("begin machine failure report: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) - txNotifications := s.newTxNotifications() if input.Stage == MachineFailureStageDaemonUninstalled { if _, err := qtx.LockMachineForLifecycle( ctx, @@ -571,8 +573,7 @@ func (s *Store) RecordMachineFailureReport( if input.Stage == MachineFailureStageDaemonUninstalled { if err := completeMachineLifecycleTerminalWorkTx( ctx, - txNotifications, - tx, + unit, qtx, input.OrgID, input.MachineID, @@ -581,5 +582,5 @@ func (s *Store) RecordMachineFailureReport( return err } } - return s.commitTxWithNotifications(ctx, tx, txNotifications, "record machine failure report") + return unit.Commit(ctx, "record machine failure report") } diff --git a/internal/storage/executionstore/machine_daemon_reconciliation_store.go b/internal/storage/executionstore/machine_daemon_reconciliation_store.go index 46268cef7..291f82d05 100644 --- a/internal/storage/executionstore/machine_daemon_reconciliation_store.go +++ b/internal/storage/executionstore/machine_daemon_reconciliation_store.go @@ -9,7 +9,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/daemonprotocol" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -17,8 +17,7 @@ import ( func reconcileRegisteredRuntimeTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input RegisterDaemonRuntimeInput, ) (DaemonRuntimeReconciliation, error) { @@ -28,7 +27,7 @@ func reconcileRegisteredRuntimeTx( } processes, err := lockAndLoadReconciliationProcesses( ctx, - tx, + unit, qtx, input.OrgID, input.MachineID, @@ -60,8 +59,7 @@ func reconcileRegisteredRuntimeTx( case ProcessStateStarting, ProcessStateRunning: if _, err := completeProcessUnknownByMachineTx( ctx, - txNotifications, - tx, + unit, qtx, input.OrgID, input.MachineID, @@ -75,8 +73,7 @@ func reconcileRegisteredRuntimeTx( ProcessStateKilled, ProcessStateUnknown: if err := completeQueuedProcessActionsForTerminalProcessTx( ctx, - txNotifications, - tx, + unit, qtx, process, ); err != nil { @@ -84,8 +81,7 @@ func reconcileRegisteredRuntimeTx( } if err := completeAcceptedProcessActionsWithoutEvidenceTx( ctx, - txNotifications, - tx, + unit, qtx, process.OrgID, process.ID, @@ -100,8 +96,7 @@ func reconcileRegisteredRuntimeTx( disposition, err := reconcileProcessTx( ctx, - txNotifications, - tx, + unit, qtx, input, process, @@ -213,11 +208,12 @@ func validateProcessReconciliationClaims( func lockAndLoadReconciliationProcesses( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, machineID uuid.UUID, claims map[uuid.UUID]ProcessReconciliationClaim, ) (map[uuid.UUID]ProcessRecord, error) { + rows, err := qtx.ListProcessesForMachineReconciliation( ctx, dbsqlc.ListProcessesForMachineReconciliationParams{ @@ -268,7 +264,7 @@ func lockAndLoadReconciliationProcesses( AgentID: record.AgentID, }) } - if err := lifecyclelock.Agents(ctx, tx, agentRefs); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { return nil, err } return loadReconciliationProcesses( @@ -332,8 +328,7 @@ func loadReconciliationProcesses( func reconcileProcessTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input RegisterDaemonRuntimeInput, process ProcessRecord, @@ -372,8 +367,7 @@ func reconcileProcessTx( case claim.ExecutionCommitted: if _, err := completeProcessUnknownByMachineTx( ctx, - txNotifications, - tx, + unit, qtx, input.OrgID, input.MachineID, @@ -389,8 +383,7 @@ func reconcileProcessTx( default: if err := completeProcessFailedWithoutExecutionTx( ctx, - txNotifications, - tx, + unit, qtx, process, ); err != nil { @@ -411,8 +404,7 @@ func reconcileProcessTx( claim.ExecutionCommitted: if _, err := completeProcessUnknownByMachineTx( ctx, - txNotifications, - tx, + unit, qtx, input.OrgID, input.MachineID, @@ -441,8 +433,7 @@ func reconcileProcessTx( claim.Phase == daemonprotocol.ProcessPhasePrepared) { if err := completeQueuedProcessActionsForTerminalProcessTx( ctx, - txNotifications, - tx, + unit, qtx, process, ); err != nil { @@ -450,8 +441,7 @@ func reconcileProcessTx( } if err := completeAcceptedProcessActionsWithoutEvidenceTx( ctx, - txNotifications, - tx, + unit, qtx, process.OrgID, process.ID, @@ -465,8 +455,7 @@ func reconcileProcessTx( } if err := completeQueuedProcessActionsForTerminalProcessTx( ctx, - txNotifications, - tx, + unit, qtx, process, ); err != nil { @@ -484,8 +473,7 @@ func reconcileProcessTx( actions, err := reconcileProcessActionsTx( ctx, - txNotifications, - tx, + unit, qtx, input, process, @@ -501,8 +489,7 @@ func reconcileProcessTx( func completeProcessFailedWithoutExecutionTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, process ProcessRecord, ) error { @@ -526,21 +513,21 @@ func completeProcessFailedWithoutExecutionTx( return fmt.Errorf("resolve process that never executed: %w", err) } failed := processRecordFromCompleteSQLC(row) - if err := completeProcessToolCallFromRecordTx( - ctx, - txNotifications, - tx, - qtx, - failed, - nil, - processExecutionNotStartedReason, - ); err != nil { + h, err := unit.Handle(failed.ProjectID, failed.AgentID) + if err != nil { return err } + if failed.ToolCallID != uuid.Nil { + if _, + err := h.CompleteProcess(ctx, + agentexecution.ProcessResult{ID: failed.ID, + Observed: nil}); err != nil { + return err + } + } if err := completeQueuedProcessActionsFailedTx( ctx, - txNotifications, - tx, + unit, qtx, process.OrgID, process.ID, @@ -550,8 +537,7 @@ func completeProcessFailedWithoutExecutionTx( } return completeAcceptedProcessActionsWithoutEvidenceTx( ctx, - txNotifications, - tx, + unit, qtx, process.OrgID, process.ID, @@ -561,8 +547,7 @@ func completeProcessFailedWithoutExecutionTx( func reconcileProcessActionsTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input RegisterDaemonRuntimeInput, process ProcessRecord, @@ -662,8 +647,7 @@ func reconcileProcessActionsTx( case locallyPresent: if err := resolveAcceptedProcessActionForReconciliationTx( ctx, - txNotifications, - tx, + unit, qtx, action, ProcessActionStateUnknown, @@ -682,8 +666,7 @@ func reconcileProcessActionsTx( default: if err := resolveAcceptedProcessActionForReconciliationTx( ctx, - txNotifications, - tx, + unit, qtx, action, ProcessActionStateFailed, @@ -753,8 +736,7 @@ func reconcileProcessActionsTx( func resolveAcceptedProcessActionForReconciliationTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, action ProcessActionRecord, state ProcessActionState, @@ -762,8 +744,7 @@ func resolveAcceptedProcessActionForReconciliationTx( ) error { _, err := completeDaemonProcessActionTx( ctx, - txNotifications, - tx, + unit, qtx, processActionCompletionInput{ ProjectID: action.ProjectID, @@ -786,9 +767,10 @@ func resolveAcceptedProcessActionForReconciliationTx( func lockProcessAgentByMachineTx( ctx context.Context, - qtx *dbsqlc.Queries, + unit *agentexecution.Unit, orgID, machineID, processID uuid.UUID, ) (ProcessRecord, bool, error) { + qtx := dbsqlc.New(unit.DB()) row, err := qtx.GetProcessByMachine( ctx, dbsqlc.GetProcessByMachineParams{OrgID: orgID, MachineID: machineID, ID: processID}, @@ -800,23 +782,23 @@ func lockProcessAgentByMachineTx( return ProcessRecord{}, false, fmt.Errorf("load process before agent lock: %w", err) } record := processRecordFromSQLC(row) - if err := lockAgentForProcessRecordTx(ctx, qtx, record); err != nil { + if err := lockAgentForProcessRecordTx(ctx, unit, record); err != nil { return ProcessRecord{}, false, err } return record, true, nil } -func lockAgentForProcessRecordTx(ctx context.Context, qtx *dbsqlc.Queries, record ProcessRecord) error { +func lockAgentForProcessRecordTx(ctx context.Context, unit *agentexecution.Unit, record ProcessRecord) error { if record.ProjectID == uuid.Nil || record.AgentID == uuid.Nil { return nil } - if _, err := qtx.LockAgentInProject( + if _, err := unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: record.ProjectID, ID: record.AgentID, }, - ); err != nil { + agentexecution.ExternalAuthority{}); err != nil { return fmt.Errorf("lock agent for process authority mutation: %w", err) } return nil diff --git a/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go b/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go index 249cdd698..81b5c82a3 100644 --- a/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go +++ b/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go @@ -41,13 +41,14 @@ func (s *Store) endExpiredDaemonRuntime( ctx context.Context, candidate dbsqlc.ListExpiredDaemonRuntimeCandidatesRow, ) (DaemonRuntimeRecord, bool, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonRuntimeRecord{}, false, fmt.Errorf("begin end expired daemon runtime: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() - txNotifications := s.newTxNotifications() qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForLifecycle( ctx, @@ -66,12 +67,7 @@ func (s *Store) endExpiredDaemonRuntime( dbsqlc.EndExpiredDaemonRuntimeParams(candidate), ) if errors.Is(err, pgx.ErrNoRows) { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "skipped expired daemon runtime", - ); err != nil { + if err := unit.Commit(ctx, "skipped expired daemon runtime"); err != nil { return DaemonRuntimeRecord{}, false, err } return DaemonRuntimeRecord{}, false, nil @@ -85,12 +81,7 @@ func (s *Store) endExpiredDaemonRuntime( record.MachineID, notifications.DaemonRuntimeEndReconnect, ) - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "end expired daemon runtime", - ); err != nil { + if err := unit.Commit(ctx, "end expired daemon runtime"); err != nil { return DaemonRuntimeRecord{}, false, err } return record, true, nil diff --git a/internal/storage/executionstore/machine_daemon_runtime_store.go b/internal/storage/executionstore/machine_daemon_runtime_store.go index 9cf76ebd5..69aa9df6f 100644 --- a/internal/storage/executionstore/machine_daemon_runtime_store.go +++ b/internal/storage/executionstore/machine_daemon_runtime_store.go @@ -56,12 +56,13 @@ func (s *Store) RegisterDaemonRuntimeWithReconciliation( } input.Capacity = normalizedJSON(input.Capacity) input.ObservedPlatform = normalizedJSON(input.ObservedPlatform) - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("begin register daemon runtime: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) machine, err := qtx.LockMachineForRuntimeRegistration( ctx, @@ -196,7 +197,7 @@ func (s *Store) RegisterDaemonRuntimeWithReconciliation( ); err != nil { return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("clear resolved machine failure report: %w", err) } - reconciliation, err := reconcileRegisteredRuntimeTx(ctx, txNotifications, tx, qtx, input) + reconciliation, err := reconcileRegisteredRuntimeTx(ctx, unit, qtx, input) if err != nil { return DaemonRuntimeRegistrationRecord{}, err } @@ -226,7 +227,7 @@ func (s *Store) RegisterDaemonRuntimeWithReconciliation( if err != nil { return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("finalize daemon runtime registration lease: %w", err) } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "register daemon runtime"); err != nil { + if err := unit.Commit(ctx, "register daemon runtime"); err != nil { return DaemonRuntimeRegistrationRecord{}, err } return DaemonRuntimeRegistrationRecord{ @@ -272,12 +273,12 @@ func (s *Store) HeartbeatDaemonRuntime( return DaemonRuntimeRecord{}, errors.New("daemon runtime lease timeout is too short") } input.ObservedPlatform = normalizedJSON(input.ObservedPlatform) - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonRuntimeRecord{}, fmt.Errorf("begin heartbeat daemon runtime: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForRuntimeRegistration( ctx, @@ -317,7 +318,7 @@ func (s *Store) HeartbeatDaemonRuntime( ); err != nil { return DaemonRuntimeRecord{}, fmt.Errorf("update machine heartbeat observation: %w", err) } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "heartbeat daemon runtime"); err != nil { + if err := unit.Commit(ctx, "heartbeat daemon runtime"); err != nil { return DaemonRuntimeRecord{}, err } return daemonRuntimeFromHeartbeat(row), nil @@ -330,12 +331,13 @@ func (s *Store) EndDaemonRuntime( if err := validateDaemonRuntimeAuthority(authority); err != nil { return DaemonRuntimeRecord{}, err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonRuntimeRecord{}, fmt.Errorf("begin end daemon runtime: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForRuntimeRegistration( ctx, @@ -368,7 +370,7 @@ func (s *Store) EndDaemonRuntime( authority.MachineID, notifications.DaemonRuntimeEndReconnect, ) - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "end daemon runtime"); err != nil { + if err := unit.Commit(ctx, "end daemon runtime"); err != nil { return DaemonRuntimeRecord{}, err } return daemonRuntimeFromEnd(row), nil @@ -381,12 +383,13 @@ func (s *Store) SleepDaemonRuntime( if err := validateDaemonRuntimeAuthority(authority); err != nil { return DaemonRuntimeRecord{}, err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonRuntimeRecord{}, fmt.Errorf("begin sleep daemon runtime: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForRuntimeRegistration( ctx, @@ -444,7 +447,7 @@ func (s *Store) SleepDaemonRuntime( authority.MachineID, notifications.DaemonRuntimeEndReconnect, ) - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "sleep daemon runtime"); err != nil { + if err := unit.Commit(ctx, "sleep daemon runtime"); err != nil { return DaemonRuntimeRecord{}, err } return daemonRuntimeFromEnd(row), nil diff --git a/internal/storage/executionstore/machine_daemon_store.go b/internal/storage/executionstore/machine_daemon_store.go index a00cb0119..46da24739 100644 --- a/internal/storage/executionstore/machine_daemon_store.go +++ b/internal/storage/executionstore/machine_daemon_store.go @@ -11,15 +11,15 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/resourcemeta" "github.com/omnara-ai/omnara/internal/resourcename" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/listing" "github.com/omnara-ai/omnara/internal/storage/storeerr" - - "github.com/omnara-ai/omnara/internal/resourcemeta" ) type CreateDaemonMachineInput struct { @@ -553,29 +553,30 @@ func (s *Store) deleteMachineOnce( ctx context.Context, input DeleteMachineInput, ) (MachineRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return MachineRecord{}, fmt.Errorf("begin delete machine: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() - machine, err := s.DeleteMachineTx(ctx, tx, txNotifications, input) + defer func() { _ = unit.Rollback(ctx) }() + + machine, err := s.DeleteMachineInUnit(ctx, unit, input) if err != nil { return MachineRecord{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "delete machine"); err != nil { + if err := unit.Commit(ctx, "delete machine"); err != nil { return MachineRecord{}, err } return machine, nil } -func (s *Store) DeleteMachineTx( +func (s *Store) DeleteMachineInUnit( ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, + unit *agentexecution.Unit, input DeleteMachineInput, ) (MachineRecord, error) { - qtx := s.q.WithTx(tx) + txNotifications := unit.Notifications() + tx := unit.DB() + qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForLifecycle( ctx, dbsqlc.LockMachineForLifecycleParams{OrgID: input.OrgID, ID: input.MachineID}, @@ -614,8 +615,7 @@ func (s *Store) DeleteMachineTx( } if err := completeMachineLifecycleTerminalWorkTx( ctx, - txNotifications, - tx, + unit, qtx, input.OrgID, input.MachineID, @@ -1008,15 +1008,16 @@ func (s *Store) deleteProjectMachineGrantOnce( ctx context.Context, orgID, projectID, id uuid.UUID, ) (ProjectMachineGrantRecord, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return ProjectMachineGrantRecord{}, fmt.Errorf( "begin delete project machine grant: %w", err, ) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) if err := lifecyclelock.EnterActiveProject(ctx, tx, orgID, projectID); err != nil { return ProjectMachineGrantRecord{}, err @@ -1045,7 +1046,7 @@ func (s *Store) deleteProjectMachineGrantOnce( if err := completeExecutionRevokedProcessesTx( ctx, txNotifications, - tx, + unit, qtx, executionRevokedProcessScope{ projectID: projectID, @@ -1067,7 +1068,7 @@ func (s *Store) deleteProjectMachineGrantOnce( return ProjectMachineGrantRecord{}, fmt.Errorf("delete project machine grant: %w", err) } record := projectMachineGrantFromDelete(row) - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "delete project machine grant"); err != nil { + if err := unit.Commit(ctx, "delete project machine grant"); err != nil { return ProjectMachineGrantRecord{}, err } return record, nil diff --git a/internal/storage/executionstore/machine_lifecycle_store.go b/internal/storage/executionstore/machine_lifecycle_store.go index f01b510e9..a1dcaa6b4 100644 --- a/internal/storage/executionstore/machine_lifecycle_store.go +++ b/internal/storage/executionstore/machine_lifecycle_store.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -188,7 +189,7 @@ func (s *Store) AdmitPoolMachineProvisioning( ) } defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + qtx := dbsqlc.New(tx) projectID, err := qtx.LockPoolMachineGrant( ctx, dbsqlc.LockPoolMachineGrantParams{ @@ -646,12 +647,13 @@ func (s *Store) ClaimExpiredIdlePoolMachineDeletion( if orgID == uuid.Nil || machineID == uuid.Nil { return PoolMachineDeletionClaim{}, false, errors.New("org and machine are required") } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return PoolMachineDeletionClaim{}, false, fmt.Errorf("begin claim expired idle pool machine deletion: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForLifecycle( ctx, @@ -678,16 +680,18 @@ func (s *Store) ClaimExpiredIdlePoolMachineDeletion( claim := expiredIdlePoolMachineDeletionClaimFromSQLC(row) claim, err = s.finalizePoolMachineDeletionClaimTx( ctx, - tx, + unit, qtx, txNotifications, claim, - "claim expired idle pool machine deletion", machineDeletingReason, ) if err != nil { return PoolMachineDeletionClaim{}, false, err } + if err := unit.Commit(ctx, "claim expired idle pool machine deletion"); err != nil { + return PoolMachineDeletionClaim{}, false, err + } return claim, true, nil } @@ -704,12 +708,13 @@ func (s *Store) ClaimPoolMachineDeletion( if input.ExpectedLifecycleVersion <= 0 { return PoolMachineDeletionClaim{}, false, errors.New("expected lifecycle version is required") } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return PoolMachineDeletionClaim{}, false, fmt.Errorf("begin claim pool machine deletion: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) if _, err := qtx.LockMachineForLifecycle( ctx, @@ -739,26 +744,27 @@ func (s *Store) ClaimPoolMachineDeletion( claim := poolMachineDeletionClaimFromSQLC(row) claim, err = s.finalizePoolMachineDeletionClaimTx( ctx, - tx, + unit, qtx, txNotifications, claim, - "claim pool machine deletion", machineDeletingReason, ) if err != nil { return PoolMachineDeletionClaim{}, false, err } + if err := unit.Commit(ctx, "claim pool machine deletion"); err != nil { + return PoolMachineDeletionClaim{}, false, err + } return claim, true, nil } func (s *Store) finalizePoolMachineDeletionClaimTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, txNotifications *notifications.TxNotifications, claim PoolMachineDeletionClaim, - operation string, terminalWorkReason string, ) (PoolMachineDeletionClaim, error) { machine := claim.Machine @@ -789,8 +795,7 @@ func (s *Store) finalizePoolMachineDeletionClaimTx( } if err := completeMachineLifecycleTerminalWorkTx( ctx, - txNotifications, - tx, + unit, qtx, machine.OrgID, machine.ID, @@ -822,9 +827,6 @@ func (s *Store) finalizePoolMachineDeletionClaimTx( if err != nil { return PoolMachineDeletionClaim{}, fmt.Errorf("finalize pool machine deletion claim: %w", err) } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, operation); err != nil { - return PoolMachineDeletionClaim{}, err - } claim.Machine.NextReconcileAfter = lease.NextReconcileAfter claim.Machine.UpdatedAt = lease.UpdatedAt return claim, nil @@ -933,12 +935,13 @@ func (s *Store) CompletePoolMachineDeletion( if deleteAttempt <= 0 { return errors.New("delete attempt is required") } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return fmt.Errorf("begin complete pool machine deletion: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) poolID, err := qtx.GetPoolMachinePoolIDForLifecycle( ctx, @@ -1003,8 +1006,7 @@ func (s *Store) CompletePoolMachineDeletion( } if err := completeMachineLifecycleTerminalWorkTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, machineID, @@ -1053,7 +1055,7 @@ func (s *Store) CompletePoolMachineDeletion( ); err != nil { return fmt.Errorf("destroy deleted-org secret versions for deletion completion: %w", err) } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "complete pool machine deletion"); err != nil { + if err := unit.Commit(ctx, "complete pool machine deletion"); err != nil { return err } return nil diff --git a/internal/storage/executionstore/machine_lifecycle_terminal_work.go b/internal/storage/executionstore/machine_lifecycle_terminal_work.go index ebdfd6acd..140ce5b28 100644 --- a/internal/storage/executionstore/machine_lifecycle_terminal_work.go +++ b/internal/storage/executionstore/machine_lifecycle_terminal_work.go @@ -2,13 +2,13 @@ package executionstore import ( "context" - "encoding/json" "errors" "fmt" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -16,8 +16,7 @@ import ( func completeMachineLifecycleTerminalWorkTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, machineID uuid.UUID, reason string, @@ -38,13 +37,12 @@ func completeMachineLifecycleTerminalWorkTx( AgentID: agent.AgentID, }) } - if err := lifecyclelock.Agents(ctx, tx, refs); err != nil { + if err := unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}); err != nil { return err } if err := completeMachineLifecycleTerminalProcessesTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, machineID, @@ -54,8 +52,7 @@ func completeMachineLifecycleTerminalWorkTx( } return completeMachineLifecycleTerminalQueuedProcessToolCallsTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, machineID, @@ -73,23 +70,28 @@ type executionRevokedProcessScope struct { func completeExecutionRevokedProcessesTx( ctx context.Context, txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, scope executionRevokedProcessScope, reason string, ) error { - lockParams := dbsqlc.LockAgentsForExecutionRevokedParams{ + lockParams := dbsqlc.ListAgentsForExecutionRevokedParams{ ProjectID: scope.projectID, AgentID: storeutil.IDFromNil(scope.agentID), ProjectMachineGrantID: storeutil.IDFromNil(scope.projectMachineGrantID), ProjectMachinePoolGrantID: storeutil.IDFromNil(scope.projectMachinePoolGrantID), } - if _, err := qtx.LockAgentsForExecutionRevoked(ctx, lockParams); err != nil { - return fmt.Errorf("lock agents for execution revoke: %w", err) + agents, err := qtx.ListAgentsForExecutionRevoked(ctx, lockParams) + if err != nil { + return fmt.Errorf("list agents for execution revoke: %w", err) + } + refs := make([]lifecyclelock.AgentRef, 0, len(agents)) + for _, agentID := range agents { + refs = append(refs, lifecyclelock.AgentRef{ProjectID: scope.projectID, AgentID: agentID}) + } + if err := unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}); err != nil { + return err } - // Binding release terminalizes its process/action work first, and new work - // requires an attached binding. These agents therefore already cover every - // process below; no second pass through the agent lock class is needed. rows, err := qtx.ListProcessesForExecutionRevoked( ctx, dbsqlc.ListProcessesForExecutionRevokedParams{ @@ -124,23 +126,23 @@ func completeExecutionRevokedProcessesTx( return fmt.Errorf("mark queued process failed for execution revoke: %w", err) } record := processRecordFromSQLC(failed) - if err := completeProcessToolCallFromRecordTx( - ctx, - txNotifications, - tx, - qtx, - record, - nil, - reason, - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return err } + if record.ToolCallID != uuid.Nil { + if _, + err := h.CompleteProcess(ctx, + agentexecution.ProcessResult{ID: record.ID, + Observed: nil}); err != nil { + return err + } + } case ProcessStateStarting, ProcessStateRunning: txNotifications.AddDaemonProcessTermination(process.MachineID, process.ID) if _, err := completeProcessUnknownByMachineTx( ctx, - txNotifications, - tx, + unit, qtx, process.OrgID, process.MachineID, @@ -153,8 +155,7 @@ func completeExecutionRevokedProcessesTx( case ProcessStateExited, ProcessStateFailed, ProcessStateKilled, ProcessStateUnknown: if err := completeUnresolvedProcessActionsForClosedProcessTx( ctx, - txNotifications, - tx, + unit, qtx, process.OrgID, process.ID, @@ -169,8 +170,7 @@ func completeExecutionRevokedProcessesTx( func completeMachineLifecycleTerminalProcessesTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, machineID uuid.UUID, reason string, @@ -190,8 +190,7 @@ func completeMachineLifecycleTerminalProcessesTx( case ProcessStateStarting, ProcessStateRunning: if _, err := completeProcessUnknownByMachineTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, machineID, @@ -205,8 +204,7 @@ func completeMachineLifecycleTerminalProcessesTx( ProcessStateKilled, ProcessStateUnknown: if err := completeUnresolvedProcessActionsForClosedProcessTx( ctx, - txNotifications, - tx, + unit, qtx, record.OrgID, record.ID, @@ -222,8 +220,7 @@ func completeMachineLifecycleTerminalProcessesTx( func completeMachineLifecycleTerminalQueuedProcessToolCallsTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, machineID uuid.UUID, reason string, @@ -245,14 +242,17 @@ func completeMachineLifecycleTerminalQueuedProcessToolCallsTx( } for _, row := range rows { process := processRecordFromSQLC(row) - failed, err := qtx.MarkQueuedProcessFailedByMachine(ctx, dbsqlc.MarkQueuedProcessFailedByMachineParams{ - ProjectID: process.ProjectID, - AgentID: process.AgentID, - ID: process.ID, - OrgID: orgID, - MachineID: machineID, - StateReasonCode: storeutil.TextFromEmpty(reason), - }) + failed, err := qtx.MarkQueuedProcessFailedByMachine( + ctx, + dbsqlc.MarkQueuedProcessFailedByMachineParams{ + ProjectID: process.ProjectID, + AgentID: process.AgentID, + ID: process.ID, + OrgID: orgID, + MachineID: machineID, + StateReasonCode: storeutil.TextFromEmpty(reason), + }, + ) if errors.Is(err, pgx.ErrNoRows) { continue } @@ -260,214 +260,30 @@ func completeMachineLifecycleTerminalQueuedProcessToolCallsTx( return fmt.Errorf("mark queued process failed for machine delete: %w", err) } record := processRecordFromSQLC(failed) - if err := completeProcessToolCallFromRecordTx(ctx, txNotifications, tx, qtx, record, nil, reason); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return err } - } - } -} - -func completeTerminalProcessActionToolCallTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - record ProcessActionRecord, - state ProcessActionState, - reason string, -) error { - applied := state == ProcessActionStateApplied - errText := reason - outcome := ToolResultOutcomeFailed - if applied { - errText = "" - outcome = ToolResultOutcomeSucceeded - } - var result json.RawMessage - if record.ActionKind == ProcessActionKindRead && !applied { - processRow, err := qtx.GetProcessForUpdate( - ctx, - dbsqlc.GetProcessForUpdateParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - ID: record.ProcessID, - }, - ) - if err != nil { - return fmt.Errorf("load process for failed read result: %w", err) - } - result, err = canonicalProcessReadFailureResult( - processRecordFromSQLC(processRow), - record, - reason, - errText, - ) - if err != nil { - return err - } - } else { - var err error - result, err = processActionToolResult( - record.ProcessID, - record.ID, - state, - reason, - errText, - ) - if err != nil { - return err - } - } - contentParts, err := ToolResultContentParts(result) - if err != nil { - return err - } - toolRow, err := qtx.CompleteToolCallFromProcessAction( - ctx, - dbsqlc.CompleteToolCallFromProcessActionParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - ToolCallID: record.ToolCallID, - ProcessID: record.ProcessID, - ProcessActionID: record.ID, - Outcome: string(outcome), - }, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - ok, checkErr := completedToolCallMissIsBenignTx( - ctx, - qtx, - record.ProjectID, - record.AgentID, - record.ToolCallID, - false, - ) - if checkErr != nil { - return checkErr - } - if ok { - return nil + if record.ToolCallID != uuid.Nil { + if _, + err := h.CompleteProcess(ctx, + agentexecution.ProcessResult{ID: record.ID, + Observed: nil}); err != nil { + return err + } } } - return fmt.Errorf("complete unknown process action tool call: %w", err) } - resultRecord := toolCallRecordFromProcessActionCompleteSQLC(toolRow) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx(ctx, txNotifications, tx, resultRecord, nil); err != nil { - return err - } - metadata, err := marshalJSON( - map[string]any{ - "reason": "process_action_" + string(state), - "process_id": record.ProcessID, - "process_action_id": record.ID, - "tool_call_id": record.ToolCallID, - }, - ) - if err != nil { - return fmt.Errorf("marshal process action wakeup metadata: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - Metadata: metadata, - }, - ); err != nil { - return fmt.Errorf("mark process action %s wakeup: %w", state, err) - } - return nil -} - -func completeProcessToolCallFromRecordTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - record ProcessRecord, - overrideResult json.RawMessage, - wakeReason string, -) error { - if record.ToolCallID == uuid.Nil { - return nil - } - outcome, result, resultErr := processToolResult(record) - if resultErr != nil { - return resultErr - } - if len(overrideResult) > 0 && string(overrideResult) != "null" { - var err error - result, err = commandTerminalToolResult(record.ID, overrideResult) - if err != nil { - return err - } - } - contentParts, err := ToolResultContentParts(result) - if err != nil { - return err - } - toolRow, err := qtx.CompleteToolCallFromProcess( - ctx, - dbsqlc.CompleteToolCallFromProcessParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - ID: record.ToolCallID, - ProcessID: record.ID, - Outcome: string(outcome), - }, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - if ok, checkErr := completedToolCallMissIsBenignTx( - ctx, - qtx, - record.ProjectID, - record.AgentID, - record.ToolCallID, - true, - ); checkErr != nil { - return checkErr - } else if ok { - return nil - } - } - return fmt.Errorf("complete process tool call: %w", err) - } - resultRecord := toolCallRecordFromProcessCompleteSQLC(toolRow) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx(ctx, txNotifications, tx, resultRecord, nil); err != nil { - return err - } - metadata, err := marshalJSON( - map[string]any{"reason": wakeReason, "process_id": record.ID, "tool_call_id": record.ToolCallID}, - ) - if err != nil { - return fmt.Errorf("marshal process tool call wakeup metadata: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - Metadata: metadata, - }, - ); err != nil { - return fmt.Errorf("mark process tool call wakeup: %w", err) - } - return nil } func completeProcessUnknownByMachineTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, machineID, processID uuid.UUID, reason, message string, ) (bool, error) { - if _, found, err := lockProcessAgentByMachineTx(ctx, qtx, orgID, machineID, processID); err != nil { + if _, found, err := lockProcessAgentByMachineTx(ctx, unit, orgID, machineID, processID); err != nil { return false, err } else if !found { return false, nil @@ -489,21 +305,21 @@ func completeProcessUnknownByMachineTx( return false, nil } record := processRecordFromSQLC(row) - if err := completeProcessToolCallFromRecordTx( - ctx, - txNotifications, - tx, - qtx, - record, - nil, - "process_unknown", - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return false, err } + if record.ToolCallID != uuid.Nil { + if _, + err := h.CompleteProcess(ctx, + agentexecution.ProcessResult{ID: record.ID, + Observed: nil}); err != nil { + return false, err + } + } if err := completeUnresolvedProcessActionsForClosedProcessTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, processID, @@ -516,16 +332,14 @@ func completeProcessUnknownByMachineTx( func completeUnresolvedProcessActionsForClosedProcessTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, processID uuid.UUID, reason string, ) error { if err := completeQueuedProcessActionsFailedTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, processID, @@ -535,8 +349,7 @@ func completeUnresolvedProcessActionsForClosedProcessTx( } return completeAcceptedProcessActionsWithoutEvidenceTx( ctx, - txNotifications, - tx, + unit, qtx, orgID, processID, @@ -546,8 +359,7 @@ func completeUnresolvedProcessActionsForClosedProcessTx( func completeQueuedProcessActionsFailedTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, processID uuid.UUID, reason string, @@ -566,25 +378,22 @@ func completeQueuedProcessActionsFailedTx( } for _, row := range rows { record := processActionRecordFromSQLC(row) - if err := completeTerminalProcessActionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - ProcessActionStateFailed, - reason, - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return err } + if record.ToolCallID != uuid.Nil { + if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return err + } + } } return nil } func completeQueuedProcessActionsForTerminalProcessTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, process ProcessRecord, ) error { @@ -600,17 +409,15 @@ func completeQueuedProcessActionsForTerminalProcessTx( } for _, row := range failedRows { record := processActionRecordFromSQLC(row) - if err := completeTerminalProcessActionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - ProcessActionStateFailed, - "process_terminal", - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return err } + if record.ToolCallID != uuid.Nil { + if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return err + } + } } resolvedRows, err := qtx.ResolveQueuedTerminateActionsForTerminalProcess( ctx, @@ -625,25 +432,22 @@ func completeQueuedProcessActionsForTerminalProcessTx( } for _, row := range resolvedRows { record := processActionRecordFromSQLC(row) - if err := completeTerminalProcessActionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - record.State, - record.StateReasonCode, - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return err } + if record.ToolCallID != uuid.Nil { + if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return err + } + } } return nil } func completeAcceptedProcessActionsWithoutEvidenceTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, orgID, processID uuid.UUID, reason string, @@ -671,17 +475,15 @@ func completeAcceptedProcessActionsWithoutEvidenceTx( } for _, row := range rows { record := processActionRecordFromSQLC(row) - if err := completeTerminalProcessActionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - record.State, - reason, - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return err } + if record.ToolCallID != uuid.Nil { + if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return err + } + } } return nil } diff --git a/internal/storage/executionstore/machine_pools_store.go b/internal/storage/executionstore/machine_pools_store.go index f55023a05..117e2a1fb 100644 --- a/internal/storage/executionstore/machine_pools_store.go +++ b/internal/storage/executionstore/machine_pools_store.go @@ -11,8 +11,9 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/resourcemeta" "github.com/omnara-ai/omnara/internal/resourcename" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/secretops" @@ -22,8 +23,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/patch" "github.com/omnara-ai/omnara/internal/storage/secretstore" "github.com/omnara-ai/omnara/internal/storage/storeerr" - - "github.com/omnara-ai/omnara/internal/resourcemeta" ) type MachineSourceKind string @@ -328,7 +327,7 @@ func (s *Store) CreateMachinePool( return MachinePoolRecord{}, fmt.Errorf("begin create machine pool: %w", err) } defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + qtx := dbsqlc.New(tx) if err := lifecyclelock.EnterActiveOrganization(ctx, tx, input.OrgID); err != nil { return MachinePoolRecord{}, err } @@ -570,7 +569,7 @@ func prepareMachinePoolConfigInput( return machinePoolDefaults{Provisioning: poolProvisioning, Environment: poolEnvironment}, nil } -func validateMachinePoolProviderAuth(ctx context.Context, tx pgx.Tx, orgID, providerAuthSecretID uuid.UUID) error { +func validateMachinePoolProviderAuth(ctx context.Context, tx dbsqlc.DBTX, orgID, providerAuthSecretID uuid.UUID) error { credential, err := secretops.LockReference(ctx, tx, orgID, providerAuthSecretID) if errors.Is(err, pgx.ErrNoRows) { return storeerr.ErrNotFound @@ -994,32 +993,34 @@ func (s *Store) DeleteMachinePool(ctx context.Context, orgID, id uuid.UUID) ([]M } func (s *Store) deleteMachinePoolOnce(ctx context.Context, orgID, id uuid.UUID) ([]MachineRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return nil, fmt.Errorf("begin delete machine pool: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - txNotifications := s.newTxNotifications() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + if err := lifecyclelock.EnterActiveOrganization(ctx, tx, orgID); err != nil { return nil, err } - machines, err := s.DeleteMachinePoolTx(ctx, tx, txNotifications, orgID, id) + machines, err := s.DeleteMachinePoolInUnit(ctx, unit, orgID, id) if err != nil { return nil, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "delete machine pool"); err != nil { + if err := unit.Commit(ctx, "delete machine pool"); err != nil { return nil, err } return machines, nil } -func (s *Store) DeleteMachinePoolTx( +func (s *Store) DeleteMachinePoolInUnit( ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, + unit *agentexecution.Unit, orgID, id uuid.UUID, ) ([]MachineRecord, error) { - qtx := s.q.WithTx(tx) + txNotifications := unit.Notifications() + tx := unit.DB() + qtx := dbsqlc.New(tx) if _, err := qtx.LockMachinePoolForUpdate( ctx, dbsqlc.LockMachinePoolForUpdateParams{OrgID: orgID, ID: id}, @@ -1074,7 +1075,7 @@ func (s *Store) DeleteMachinePoolTx( AgentID: agentRow.AgentID, }) } - if err := lifecyclelock.Agents(ctx, tx, agentRefs); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { return nil, err } if _, err := qtx.DeleteMachinePool( @@ -1097,7 +1098,7 @@ func (s *Store) DeleteMachinePoolTx( if err := completeExecutionRevokedProcessesTx( ctx, txNotifications, - tx, + unit, qtx, executionRevokedProcessScope{ projectID: poolGrantRef.ProjectID, diff --git a/internal/storage/executionstore/model_call_schema_guards_integration_test.go b/internal/storage/executionstore/model_call_schema_guards_integration_test.go index e4975c185..3afbf7bac 100644 --- a/internal/storage/executionstore/model_call_schema_guards_integration_test.go +++ b/internal/storage/executionstore/model_call_schema_guards_integration_test.go @@ -14,23 +14,22 @@ import ( "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/stretchr/testify/require" ) func TestConcurrentNormalModelCallClaimsHaveOneSender(t *testing.T) { t.Parallel() ctx := context.Background() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "normal_creator_authority") + fixture, admitted, _ := newMultiInputContinuationSeedFixture(t, ctx, "normal_creator_authority") openingInputIDs := make([]uuid.UUID, 0, len(admitted.Inputs)) for _, input := range admitted.Inputs { openingInputIDs = append(openingInputIDs, input.ID) } - claimInput := executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[len(admitted.Events)-1].Sequence, + claimInput := executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, } type result struct { @@ -45,7 +44,8 @@ func TestConcurrentNormalModelCallClaimsHaveOneSender(t *testing.T) { go func() { defer workers.Done() <-start - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, claimInput) + prepared1, err := fixture.Store.Execution().PrepareNormalModelCall(ctx, claimInput) + claim := prepared1.Claim results <- result{claim: claim, err: err} }() } @@ -68,50 +68,18 @@ func TestConcurrentNormalModelCallClaimsHaveOneSender(t *testing.T) { } func TestModelCallOpeningInputsRejectWrongProject(t *testing.T) { - t.Parallel() - ctx := context.Background() - fixture, admitted, _ := newMultiInputContinuationSeedFixture( - t, - ctx, - "opening_inputs_wrong_project", - ) - wrongProjectID := seedAdditionalProjectForTest( - t, - ctx, - fixture.Store.pool, - "opening_inputs_wrong_scope", - ) - frontier := admitted.Events[len(admitted.Events)-1].Sequence - - var correctProjectCount, wrongProjectCount int - if err := fixture.Store.pool.QueryRow( - ctx, - `SELECT count(*) FROM agent_model_call_opening_content_inputs($1, $2, $3, $4)`, - testProjectID, - fixture.AgentID, - admitted.Turn.ID, - frontier, - ).Scan(&correctProjectCount); err != nil { - t.Fatalf("load correct-project opening inputs: %v", err) - } - if err := fixture.Store.pool.QueryRow( - ctx, - `SELECT count(*) FROM agent_model_call_opening_content_inputs($1, $2, $3, $4)`, - wrongProjectID, - fixture.AgentID, - admitted.Turn.ID, - frontier, - ).Scan(&wrongProjectCount); err != nil { - t.Fatalf("load wrong-project opening inputs: %v", err) - } - if correctProjectCount != len(admitted.Inputs) || wrongProjectCount != 0 { - t.Fatalf( - "opening input counts correct=%d wrong=%d, want %d/0", - correctProjectCount, - wrongProjectCount, - len(admitted.Inputs), - ) - } + ctx := t.Context() + f, _, claim := newStartedNormalModelCallTestFixture(t, ctx, "opening-project-scope") + _, found, err := f.Store.Execution().GetModelCallContext(ctx, uuid.New(), f.AgentID, claim.Context.ID) + require.NoError(t, err) + require.False(t, found) + _, err = f.Store.Execution(). + PrepareNormalModelCall(ctx, + executionstore.PrepareNormalModelCallInput{ProjectID: uuid.New(), + AgentID: f.AgentID, + RuntimeLockID: f.Lock.ID, + OpeningInputIDs: []uuid.UUID{uuid.New()}}) + require.Error(t, err) } func TestModelCallContextDatabaseGuardsRejectRebinding(t *testing.T) { @@ -123,8 +91,15 @@ func TestModelCallContextDatabaseGuardsRejectRebinding(t *testing.T) { query string secondArg bool }{ - {name: "attempt number", query: `UPDATE model_call_contexts SET attempt_number = attempt_number + 1 WHERE id = $1`}, - {name: "runtime lock", query: `UPDATE model_call_contexts SET runtime_lock_id = $2 WHERE id = $1`, secondArg: true}, + { + name: "attempt number", + query: `UPDATE model_call_contexts SET attempt_number = attempt_number + 1 WHERE id = $1`, + }, + { + name: "runtime lock", + query: `UPDATE model_call_contexts SET runtime_lock_id = $2 WHERE id = $1`, + secondArg: true, + }, { name: "event frontier", query: `UPDATE model_call_contexts SET input_event_sequence = input_event_sequence + 1 WHERE id = $1`, @@ -187,12 +162,13 @@ func TestModelCallContextIdentityIndexesRejectDuplicateLogicalContexts(t *testin normalFixture, _, normal := newStartedNormalModelCallTestFixture(t, ctx, "normal_identity_guard") _, err := normalFixture.Store.pool.Exec(ctx, ` INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, + turn_id, opening_input_ids, opening_event_sequence, +org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, runtime_lock_id, state, created_at ) -SELECT org_id, project_id, agent_id, operation_kind, +SELECT turn_id, opening_input_ids, opening_event_sequence, org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, runtime_lock_id, 'started', created_at + interval '1 second' @@ -200,7 +176,11 @@ FROM model_call_contexts WHERE id = $1`, normal.Context.ID) assertPgConstraint(t, err, "23505", "model_call_contexts_normal_identity_idx") - compactionFixture, admitted, _ := newMultiInputContinuationSeedFixture(t, ctx, "compaction_identity_guard") + compactionFixture, admitted, _ := newMultiInputContinuationSeedFixture( + t, + ctx, + "compaction_identity_guard", + ) frontier := admitted.Events[len(admitted.Events)-1].Sequence compaction := claimSentCompactionForRangeTest( t, @@ -213,13 +193,14 @@ WHERE id = $1`, normal.Context.ID) ) _, err = compactionFixture.Store.pool.Exec(ctx, ` INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, + turn_id, opening_input_ids, opening_event_sequence, +org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, source_event_sequence_end, runtime_lock_id, state, created_at ) -SELECT org_id, project_id, agent_id, operation_kind, +SELECT turn_id, opening_input_ids, opening_event_sequence, org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, source_event_sequence_end, @@ -316,8 +297,8 @@ func TestModelCallOperationRejectsSecondSemanticOutcome(t *testing.T) { var turnID uuid.UUID if err := fixture.Store.pool.QueryRow(ctx, ` SELECT turn_id -FROM model_call_context_turns -WHERE project_id = $1 AND agent_id = $2 AND model_call_context_id = $3 +FROM model_call_contexts +WHERE project_id = $1 AND agent_id = $2 AND id = $3 `, testProjectID, fixture.AgentID, first.Context.ID).Scan(&turnID); err != nil { t.Fatalf("load model call turn: %v", err) } @@ -336,12 +317,13 @@ WHERE project_id = $1 AND agent_id = $2 AND model_call_context_id = $3 var secondContextID uuid.UUID if err := fixture.Store.pool.QueryRow(ctx, ` INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, + turn_id, opening_input_ids, opening_event_sequence, +org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, runtime_lock_id, state, created_at ) -SELECT org_id, project_id, agent_id, operation_kind, +SELECT turn_id, opening_input_ids, opening_event_sequence, org_id, project_id, agent_id, operation_kind, attempt_number + 1, agent_config_id, configured_model_revision_id, input_event_sequence, runtime_lock_id, 'started', statement_timestamp() diff --git a/internal/storage/executionstore/model_context_data.go b/internal/storage/executionstore/model_context_data.go new file mode 100644 index 000000000..7b66f4ed5 --- /dev/null +++ b/internal/storage/executionstore/model_context_data.go @@ -0,0 +1,83 @@ +package executionstore + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" +) + +type ModelContextData struct { + ProjectID uuid.UUID + AgentID uuid.UUID + InputEventSequence int64 + Checkpoint *ContextCheckpointRecord + CheckpointEndsWithOutputLimit bool + Events []ContextEventRecord + ToolCalls []ToolCallRecord + ProviderReplayCutoffEventSequence int64 +} + +func loadModelContextDataTx( + ctx context.Context, q *dbsqlc.Queries, model ModelCallContextRecord, +) (*ModelContextData, error) { + result := &ModelContextData{ + ProjectID: model.ProjectID, AgentID: model.AgentID, InputEventSequence: model.InputEventSequence, + } + checkpoint, err := q.GetLatestApplicableContextCheckpoint(ctx, dbsqlc.GetLatestApplicableContextCheckpointParams{ + ProjectID: model.ProjectID, AgentID: model.AgentID, MaxEventSequence: model.InputEventSequence, + }) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return nil, fmt.Errorf("load model checkpoint: %w", err) + } + after := int64(0) + if err == nil { + record := contextCheckpointRecordFromLatestSQLC(checkpoint) + result.Checkpoint = &record + after = record.SummarizedThroughEventSequence + result.CheckpointEndsWithOutputLimit, err = q.IsOutputLimitBoundary(ctx, dbsqlc.IsOutputLimitBoundaryParams{ + ProjectID: model.ProjectID, AgentID: model.AgentID, EventSequence: after, + }) + if err != nil { + return nil, fmt.Errorf("load checkpoint boundary: %w", err) + } + } + for cursor := after; ; { + rows, err := q.ListContextEvents(ctx, dbsqlc.ListContextEventsParams{ + ProjectID: model.ProjectID, AgentID: model.AgentID, + AfterSequence: cursor, Watermark: model.InputEventSequence, PageLimit: 500, + }) + if err != nil { + return nil, fmt.Errorf("load model transcript: %w", err) + } + result.Events = append(result.Events, contextEventsFromSQLC(rows, model.ProjectID, model.AgentID)...) + if len(rows) < 500 { + break + } + cursor = rows[len(rows)-1].Sequence + if cursor >= model.InputEventSequence { + break + } + } + tools, err := q.ListCompletedToolCallsAtWatermark(ctx, dbsqlc.ListCompletedToolCallsAtWatermarkParams{ + ProjectID: model.ProjectID, AgentID: model.AgentID, + AfterEventSequence: after, MaxEventSequence: model.InputEventSequence, + }) + if err != nil { + return nil, fmt.Errorf("load model tool results: %w", err) + } + for _, row := range tools { + result.ToolCalls = append(result.ToolCalls, toolCallRecordFromWatermarkSQLC(row)) + } + result.ProviderReplayCutoffEventSequence, err = q.GetProviderReplaySuppressionCutoff(ctx, + dbsqlc.GetProviderReplaySuppressionCutoffParams{ + ProjectID: model.ProjectID, AgentID: model.AgentID, ModelCallContextID: model.ID, + }) + if err != nil { + return nil, fmt.Errorf("load model replay policy: %w", err) + } + return result, nil +} diff --git a/internal/storage/executionstore/model_context_data_integration_test.go b/internal/storage/executionstore/model_context_data_integration_test.go new file mode 100644 index 000000000..cdcc9b831 --- /dev/null +++ b/internal/storage/executionstore/model_context_data_integration_test.go @@ -0,0 +1,43 @@ +//go:build integration + +package executionstore_test + +import ( + "testing" + + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/stretchr/testify/require" +) + +func TestPreparedModelContextDataMatchesWatermarkReads(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "prepared_data") + ids := createToolCallBatchForProcessTest(t, ctx, f, "prepared_data", []processToolCallBatchItem{ + builtInProcessToolCallBatchItem("first", "read_process"), + builtInProcessToolCallBatchItem("second", "read_process"), + }) + for _, id := range ids { + ownedCompleteTool(t, ctx, f, id) + } + advance := ownedAdvanceInput(f, true) + advance.PrepareModel = true + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, advance) + require.NoError(t, err) + require.True(t, continued) + require.Equal(t, executionstore.ModelWorkContinue, work.Model.Kind) + data := work.Model.Prepared.ContextData + require.NotNil(t, data) + events, err := f.Store.Execution().ListContextEvents(ctx, testProjectID, f.AgentID, 0, data.InputEventSequence, 500) + require.NoError(t, err) + require.Equal(t, events, data.Events) + calls, err := f.Store.Execution().ListCompletedToolCallsAtWatermark( + ctx, testProjectID, f.AgentID, 0, data.InputEventSequence, + ) + require.NoError(t, err) + require.Len(t, calls, 2) + require.Equal(t, calls, data.ToolCalls) + ownedInput(t, ctx, f, executionstore.DeliveryModeSteering, "later") + frozen, err := f.Store.Execution().ListContextEvents(ctx, testProjectID, f.AgentID, 0, data.InputEventSequence, 500) + require.NoError(t, err) + require.Equal(t, events, frozen) +} diff --git a/internal/storage/executionstore/model_contexts_boundary.go b/internal/storage/executionstore/model_contexts_boundary.go deleted file mode 100644 index 45cfda726..000000000 --- a/internal/storage/executionstore/model_contexts_boundary.go +++ /dev/null @@ -1,86 +0,0 @@ -package executionstore - -import ( - "context" - "fmt" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" -) - -func resolveModelReadyBoundaryTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - q *dbsqlc.Queries, - projectID, agentID, modelCallContextID uuid.UUID, -) (bool, error) { - hasLaterSemanticEvent, err := q.ModelCallContextHasLaterSemanticEvent( - ctx, - dbsqlc.ModelCallContextHasLaterSemanticEventParams{ - ProjectID: projectID, - AgentID: agentID, - ModelCallContextID: modelCallContextID, - }, - ) - if err != nil { - return false, fmt.Errorf("check model call context semantic frontier: %w", err) - } - steering, err := selectLockedSteeringAgentInputsForAdmissionTx(ctx, q, projectID, agentID) - if err != nil { - return false, err - } - if len(steering) == 0 { - return hasLaterSemanticEvent, nil - } - if _, err := admitLockedAgentInputsAndOpenTurnTx( - ctx, - txNotifications, - tx, - q, - admitAgentInputAndOpenTurnInput{ - ProjectID: projectID, - AgentID: agentID, - }, - steering, - ); err != nil { - return false, fmt.Errorf("admit steering at model-ready boundary: %w", err) - } - return true, nil -} - -func supersedeUnacceptedAtLaterModelReadyBoundaryTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - q *dbsqlc.Queries, - projectID, agentID, modelCallContextID uuid.UUID, - wakeupReason string, -) (bool, error) { - preempted, err := resolveModelReadyBoundaryTx( - ctx, - txNotifications, - tx, - q, - projectID, - agentID, - modelCallContextID, - ) - if err != nil || !preempted { - return preempted, err - } - metadata, err := marshalJSON(map[string]string{"reason": wakeupReason}) - if err != nil { - return false, fmt.Errorf("marshal model-ready boundary wakeup metadata: %w", err) - } - if err := q.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: projectID, - AgentID: agentID, - Metadata: metadata, - }); err != nil { - return false, fmt.Errorf("reconcile wakeup after model-ready boundary: %w", err) - } - return true, nil -} diff --git a/internal/storage/executionstore/model_contexts_claims.go b/internal/storage/executionstore/model_contexts_claims.go index 543f62d33..0c88910bc 100644 --- a/internal/storage/executionstore/model_contexts_claims.go +++ b/internal/storage/executionstore/model_contexts_claims.go @@ -4,593 +4,73 @@ import ( "context" "errors" "fmt" - "slices" "github.com/google/uuid" "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/modelprotocol" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) -func (s *Store) ClaimNormalModelCall( - ctx context.Context, - input ClaimNormalModelCallInput, -) (ModelCallClaim, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - len(input.OpeningInputIDs) == 0 || input.AgentConfigID == uuid.Nil || - input.InputEventSequence <= 0 { - return ModelCallClaim{}, errors.New( - "project, agent, runtime, opening inputs, agent config, and input event sequence are required", - ) - } - if (input.SourceModelCallContextID == uuid.Nil) != (input.SourceModelOutputID == uuid.Nil) { - return ModelCallClaim{}, errors.New( - "continuation source model context and output must be provided together", - ) - } - return s.claimModelCall(ctx, claimModelCallInput{normal: &input}) -} - func (s *Store) ClaimCompactionModelCall( ctx context.Context, input ClaimCompactionModelCallInput, ) (ModelCallClaim, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.InputEventSequence <= 0 || input.SourceEventSequenceEnd <= 0 || - input.SourceEventSequenceEnd > input.InputEventSequence || - input.ParentContextID == uuid.Nil { - return ModelCallClaim{}, errors.New( - "project, agent, runtime, parent context, frontier, and a valid compaction source range are required", - ) - } - return s.claimModelCall(ctx, claimModelCallInput{compaction: &input}) -} - -type claimModelCallInput struct { - normal *ClaimNormalModelCallInput - compaction *ClaimCompactionModelCallInput -} - -func (s *Store) claimModelCall(ctx context.Context, input claimModelCallInput) (ModelCallClaim, error) { - projectID, agentID, runtimeLockID := claimIdentity(input) - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return ModelCallClaim{}, fmt.Errorf("begin claim model call: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - - q := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - projectID, - agentID, - runtimeLockID, - ); err != nil { - return ModelCallClaim{}, err - } - - var claim modelCallContextClaimTx - if input.normal != nil { - claim, err = claimNormalContextTx(ctx, q, *input.normal) - } else { - claim, err = claimCompactionContextTx(ctx, q, *input.compaction) - } - if err != nil { - return ModelCallClaim{}, err - } - result, err := applyModelCallAdmissionTx( - ctx, - txNotifications, - tx, - q, - claim, - runtimeLockID, - ) - if err != nil { - return ModelCallClaim{}, err - } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "claim model call"); err != nil { - return ModelCallClaim{}, err - } - return result, nil -} - -func applyModelCallAdmissionTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - q *dbsqlc.Queries, - claim modelCallContextClaimTx, - runtimeLockID uuid.UUID, -) (ModelCallClaim, error) { - if !claim.created { - return ModelCallClaim{Context: claim.context}, nil - } - if claim.context.State != ModelCallContextStarted { - return ModelCallClaim{}, errors.New("new model call context was not started") - } - if claim.context.RuntimeLockID != runtimeLockID { - return ModelCallClaim{}, errors.New("new model call context was not runtime-owned") - } - if claim.newManagedWorkAllowed { - return ModelCallClaim{ - Context: claim.context, - Created: true, - Claimed: true, - }, nil - } - failure, err := recordTerminalModelCallFailureTx( - ctx, - txNotifications, - tx, - q, - managedWorkAdmissionModelFailure(claim.context), - modelCallContextRuntimeOwned, - claim.context.OperationKind, - ) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { return ModelCallClaim{}, err } - return ModelCallClaim{Context: failure.context, Created: true}, nil -} - -func managedWorkAdmissionModelFailure( - contextRow ModelCallContextRecord, -) RecordModelCallErrorAndCompleteContextInput { - return RecordModelCallErrorAndCompleteContextInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - RuntimeLockID: contextRow.RuntimeLockID, - ModelCallContextID: contextRow.ID, - ErrorKind: modelprotocol.ErrorKindRuntime, - ErrorCode: storeerr.ManagedWorkAdmissionDeniedCode, - ErrorMessage: storeerr.InsufficientOmnaraCreditsMessage, - } -} - -func claimIdentity(input claimModelCallInput) (uuid.UUID, uuid.UUID, uuid.UUID) { - if input.normal != nil { - return input.normal.ProjectID, input.normal.AgentID, input.normal.RuntimeLockID - } - return input.compaction.ProjectID, input.compaction.AgentID, input.compaction.RuntimeLockID -} - -func claimNormalContextTx( - ctx context.Context, - q *dbsqlc.Queries, - input ClaimNormalModelCallInput, -) (modelCallContextClaimTx, error) { - latestEventSequence, err := q.MaxEventSequence(ctx, dbsqlc.MaxEventSequenceParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }) - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("load model call event frontier: %w", err) - } - if latestEventSequence != input.InputEventSequence { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - - matches, err := q.OpeningContentInputSetMatchesInputSequence( - ctx, - dbsqlc.OpeningContentInputSetMatchesInputSequenceParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - InputEventSequence: input.InputEventSequence, - InputIds: input.OpeningInputIDs, - }, - ) - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("validate model call opening inputs: %w", err) - } - if !matches { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - - identity := dbsqlc.GetNormalModelCallContextByIdentityParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - InputEventSequence: input.InputEventSequence, - } - id, err := q.GetNormalModelCallContextByIdentity(ctx, identity) - if err == nil { - row, loadErr := loadModelCallContextByID(ctx, q, input.ProjectID, input.AgentID, id) - if loadErr == nil && row.AgentConfigID != input.AgentConfigID { - return modelCallContextClaimTx{}, fmt.Errorf( - "agent config changed after model call context creation: %w", - storeerr.ErrStateTransitionConflict, - ) - } - return modelCallContextClaimTx{context: row}, loadErr - } - if !errors.Is(err, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, fmt.Errorf("load normal model call context: %w", err) - } - - work, err := q.NextAgentModelWork( - ctx, - dbsqlc.NextAgentModelWorkParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }, - ) - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, fmt.Errorf("validate selected model work: %w", err) - } - if errors.Is(err, pgx.ErrNoRows) || !slices.Equal(work.InputIds, input.OpeningInputIDs) { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - if input.SourceModelCallContextID == uuid.Nil { - if ModelWorkKind(work.WorkKind) != ModelWorkStart { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - } else if ModelWorkKind(work.WorkKind) != ModelWorkContinue || - work.ModelCallContextID != input.SourceModelCallContextID || - work.ModelOutputID != input.SourceModelOutputID { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - - if live, err := q.AgentHasLiveModelCallContextBeforeFrontier( - ctx, - dbsqlc.AgentHasLiveModelCallContextBeforeFrontierParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - InputEventSequence: input.InputEventSequence, - }, - ); err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("check older live model call contexts: %w", err) - } else if live { - return modelCallContextClaimTx{}, storeerr.ErrStateTransitionConflict - } - modelRevision, err := getModelCallRevisionForClaim( - ctx, - q, - input.ProjectID, - input.AgentConfigID, - ) - if err != nil { - return modelCallContextClaimTx{}, err - } - - id, err = q.InsertNormalModelCallContext(ctx, dbsqlc.InsertNormalModelCallContextParams{ - InputEventSequence: input.InputEventSequence, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - RuntimeLockID: input.RuntimeLockID, - ConfiguredModelRevisionID: modelRevision.ID, + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.ResumeCompaction(ctx, agentexecution.ResumeCompactionInput{ + RuntimeLockID: input.RuntimeLockID, ParentID: input.ParentContextID, + Watermark: input.InputEventSequence, SourceEnd: input.SourceEventSequenceEnd, }) - if errors.Is(err, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("create normal model call context: %w", err) - } - row, err := loadModelCallContextByID(ctx, q, input.ProjectID, input.AgentID, id) if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("load normal model call context: %w", err) - } - if row.AgentConfigID != input.AgentConfigID { - return modelCallContextClaimTx{}, fmt.Errorf( - "agent config changed before model call context creation: %w", - storeerr.ErrStateTransitionConflict, - ) + return ModelCallClaim{}, err } - return modelCallContextClaimTx{ - context: row, - created: true, - newManagedWorkAllowed: modelRevision.NewManagedWorkAllowed, - }, nil -} - -type modelCallContextClaimTx struct { - context ModelCallContextRecord - created bool - newManagedWorkAllowed bool -} - -func claimCompactionContextTx( - ctx context.Context, - q *dbsqlc.Queries, - input ClaimCompactionModelCallInput, -) (modelCallContextClaimTx, error) { - parent, err := loadModelCallContextByID( - ctx, - q, + if err := applyAdmissionDestination(ctx, + unit, input.ProjectID, input.AgentID, - input.ParentContextID, - ) - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("load parent model call context: %w", err) - } - if parent.OperationKind != ModelCallOperationNormal || - parent.InputEventSequence != input.InputEventSequence || - parent.State != ModelCallContextFailed || - parent.RecoveryKind != ModelCallRecoveryCompact { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - summarizedThrough := int64(0) - checkpoint, checkpointErr := q.GetLatestApplicableContextCheckpoint( - ctx, - dbsqlc.GetLatestApplicableContextCheckpointParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - MaxEventSequence: input.InputEventSequence, - }, - ) - if checkpointErr == nil { - summarizedThrough = checkpoint.SummarizedThroughEventSequence - } else if !errors.Is(checkpointErr, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, fmt.Errorf("load prior compaction checkpoint: %w", checkpointErr) - } - if input.SourceEventSequenceEnd <= summarizedThrough { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable - } - sourceEnd := input.SourceEventSequenceEnd - identity := dbsqlc.GetCompactionModelCallContextByIdentityParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - InputEventSequence: input.InputEventSequence, - SourceEventSequenceEnd: &sourceEnd, - } - id, err := q.GetCompactionModelCallContextByIdentity(ctx, identity) - if err == nil { - row, loadErr := loadModelCallContextByID(ctx, q, input.ProjectID, input.AgentID, id) - if loadErr == nil && row.AgentConfigID != parent.AgentConfigID { - return modelCallContextClaimTx{}, fmt.Errorf( - "agent config changed after compaction context creation: %w", - storeerr.ErrStateTransitionConflict, - ) - } - return modelCallContextClaimTx{context: row}, loadErr - } - if !errors.Is(err, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, fmt.Errorf("load compaction model call context: %w", err) - } - modelRevision, err := getModelCallRevisionForClaim( - ctx, - q, - input.ProjectID, - parent.AgentConfigID, - ) - if err != nil { - return modelCallContextClaimTx{}, err - } - - id, err = q.InsertTriggeredCompactionModelCallContext( - ctx, - dbsqlc.InsertTriggeredCompactionModelCallContextParams{ - SourceEventSequenceEnd: &sourceEnd, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - RuntimeLockID: input.RuntimeLockID, - ParentModelCallContextID: input.ParentContextID, - ConfiguredModelRevisionID: modelRevision.ID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, storeerr.ErrAgentNotAdvanceable + result.Admission); err != nil { + return ModelCallClaim{}, err } - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("create compaction model call context: %w", err) + if result.Preempted { + return ModelCallClaim{}, storeerr.ErrAgentNotAdvanceable } - row, err := loadModelCallContextByID(ctx, q, input.ProjectID, input.AgentID, id) + claim, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, result.Model) if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf("load compaction model call context: %w", err) + return ModelCallClaim{}, err } - if row.AgentConfigID != parent.AgentConfigID { - return modelCallContextClaimTx{}, fmt.Errorf( - "agent config changed before compaction context creation: %w", - storeerr.ErrStateTransitionConflict, - ) + if err = unit.Commit(ctx, "claim compaction"); err != nil { + return ModelCallClaim{}, err } - return modelCallContextClaimTx{ - context: row, - created: true, - newManagedWorkAllowed: modelRevision.NewManagedWorkAllowed, - }, nil + return claim, nil } func (s *Store) ClaimNextModelCallContext( ctx context.Context, input ClaimNextModelCallContextInput, ) (ModelCallClaim, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || - input.PredecessorModelCallContextID == uuid.Nil || input.RuntimeLockID == uuid.Nil { - return ModelCallClaim{}, errors.New("project, agent, predecessor context, and runtime are required") - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return ModelCallClaim{}, fmt.Errorf("begin claim next model call context: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return ModelCallClaim{}, err } - predecessor, err := loadModelCallContextByID( - ctx, - q, - input.ProjectID, - input.AgentID, - input.PredecessorModelCallContextID, - ) - if err != nil { - return ModelCallClaim{}, fmt.Errorf("load model call context for retry: %w", err) - } - claim, err := claimNextModelCallContextTx( - ctx, - q, - predecessor, - input.RuntimeLockID, - ) + defer func() { _ = unit.Rollback(ctx) }() + prepared, err := h.ResumeModel(ctx, input.RuntimeLockID, input.PredecessorModelCallContextID) if err != nil { return ModelCallClaim{}, err } - result, err := applyModelCallAdmissionTx( - ctx, - txNotifications, - tx, - q, - claim, - input.RuntimeLockID, - ) + result, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, prepared) if err != nil { return ModelCallClaim{}, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "claim next model call context", - ); err != nil { + if err = unit.Commit(ctx, "claim retry"); err != nil { return ModelCallClaim{}, err } return result, nil } -func claimNextModelCallContextTx( - ctx context.Context, - q *dbsqlc.Queries, - predecessor ModelCallContextRecord, - runtimeLockID uuid.UUID, -) (modelCallContextClaimTx, error) { - latest, err := loadLatestModelCallContextForOperation(ctx, q, predecessor) - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf( - "load latest model call context for operation: %w", - err, - ) - } - if latest.ID != predecessor.ID { - return modelCallContextClaimTx{context: latest}, nil - } - if predecessor.State == ModelCallContextFailed && - predecessor.RecoveryKind == ModelCallRecoveryRetry { - modelRevision, err := getModelCallRevisionForClaim( - ctx, - q, - predecessor.ProjectID, - predecessor.AgentConfigID, - ) - if err != nil { - return modelCallContextClaimTx{}, err - } - id, err := q.InsertNextModelCallContext(ctx, dbsqlc.InsertNextModelCallContextParams{ - ProjectID: predecessor.ProjectID, - AgentID: predecessor.AgentID, - PredecessorModelCallContextID: predecessor.ID, - MaxRetries: MaxModelCallRetriesPerOperation, - ConfiguredModelRevisionID: modelRevision.ID, - RuntimeLockID: runtimeLockID, - }) - if err == nil { - row, loadErr := loadModelCallContextByID( - ctx, - q, - predecessor.ProjectID, - predecessor.AgentID, - id, - ) - return modelCallContextClaimTx{ - context: row, - created: loadErr == nil, - newManagedWorkAllowed: modelRevision.NewManagedWorkAllowed, - }, loadErr - } - if !errors.Is(err, pgx.ErrNoRows) { - return modelCallContextClaimTx{}, fmt.Errorf("claim next model call context: %w", err) - } - latest, err = loadLatestModelCallContextForOperation(ctx, q, predecessor) - if err != nil { - return modelCallContextClaimTx{}, fmt.Errorf( - "reload latest model call context for operation: %w", - err, - ) - } - } - return modelCallContextClaimTx{context: latest}, nil -} - -type modelCallRevisionForClaim struct { - ID uuid.UUID - NewManagedWorkAllowed bool -} - -func getModelCallRevisionForClaim( - ctx context.Context, - q *dbsqlc.Queries, - projectID, agentConfigID uuid.UUID, -) (modelCallRevisionForClaim, error) { - revision, err := q.GetModelCallRevisionForClaim( - ctx, - dbsqlc.GetModelCallRevisionForClaimParams{ - ProjectID: projectID, - AgentConfigID: agentConfigID, - }, - ) - if err != nil { - return modelCallRevisionForClaim{}, fmt.Errorf( - "resolve configured model revision for model call: %w", - err, - ) - } - return modelCallRevisionForClaim{ - ID: revision.CurrentRevisionID, - NewManagedWorkAllowed: revision.NewManagedWorkAllowed, - }, nil -} - -func loadLatestModelCallContextForOperation( - ctx context.Context, - q *dbsqlc.Queries, - contextRow ModelCallContextRecord, -) (ModelCallContextRecord, error) { - var id uuid.UUID - var err error - switch contextRow.OperationKind { - case ModelCallOperationNormal: - id, err = q.GetNormalModelCallContextByIdentity( - ctx, - dbsqlc.GetNormalModelCallContextByIdentityParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - InputEventSequence: contextRow.InputEventSequence, - }, - ) - case ModelCallOperationCompaction: - id, err = q.GetCompactionModelCallContextByIdentity( - ctx, - dbsqlc.GetCompactionModelCallContextByIdentityParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - InputEventSequence: contextRow.InputEventSequence, - SourceEventSequenceEnd: contextRow.SourceEventSequenceEnd, - }, - ) - default: - return ModelCallContextRecord{}, fmt.Errorf( - "unsupported model call operation %q", - contextRow.OperationKind, - ) - } - if err != nil { - return ModelCallContextRecord{}, err - } - return loadModelCallContextByID(ctx, q, contextRow.ProjectID, contextRow.AgentID, id) -} - func (s *Store) GetModelCallContext( ctx context.Context, projectID, agentID, id uuid.UUID, @@ -630,11 +110,17 @@ func (s *Store) GetNormalModelCallContextForFrontier( return ModelCallContextRecord{}, false, nil } if err != nil { - return ModelCallContextRecord{}, false, fmt.Errorf("get normal model call context for frontier: %w", err) + return ModelCallContextRecord{}, false, fmt.Errorf( + "get normal model call context for frontier: %w", + err, + ) } record, err := loadModelCallContextByID(ctx, s.q, projectID, agentID, id) if err != nil { - return ModelCallContextRecord{}, false, fmt.Errorf("load normal model call context for frontier: %w", err) + return ModelCallContextRecord{}, false, fmt.Errorf( + "load normal model call context for frontier: %w", + err, + ) } return record, true, nil } @@ -657,7 +143,7 @@ func loadModelCallContextByID( func loadModelCallContextByIDTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID, id uuid.UUID, ) (ModelCallContextRecord, error) { return loadModelCallContextByID(ctx, dbsqlc.New(tx), projectID, agentID, id) diff --git a/internal/storage/executionstore/model_contexts_compaction.go b/internal/storage/executionstore/model_contexts_compaction.go index 97296e4e5..abf9c6cc6 100644 --- a/internal/storage/executionstore/model_contexts_compaction.go +++ b/internal/storage/executionstore/model_contexts_compaction.go @@ -2,13 +2,9 @@ package executionstore import ( "context" - "encoding/json" - "errors" - "fmt" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -17,192 +13,53 @@ func (s *Store) RecordModelCallFailureAndClaimCompaction( ctx context.Context, input RecordModelCallFailureAndClaimCompactionInput, ) (TriggeredCompactionHandoff, error) { - failure := input.Failure - if input.ParentContextID == uuid.Nil || input.SourceEventSequenceEnd <= 0 { - return TriggeredCompactionHandoff{}, errors.New( - "parent model context and a valid compaction source range are required", - ) - } - if failure.RecoveryKind != ModelCallRecoveryCompact { - return TriggeredCompactionHandoff{}, errors.New("compaction recovery is required") - } - if err := validateRecoverableModelCallFailure(failure); err != nil { - return TriggeredCompactionHandoff{}, err - } - var err error - failure.ErrorDetails, err = normalizedJSONObject(failure.ErrorDetails, "model call error details") - if err != nil { - return TriggeredCompactionHandoff{}, err - } - if err := validateModelCallFailureEvidence( - failure.APIFormat, - failure.APIVariant, - "", - failure.ProviderRequestID, - failure.ProviderResponseID, - modelenvelope.NormalizeUsage(failure.Usage) != (modelenvelope.Usage{}), - failure.ProviderReportedCostUSD, - ); err != nil { - return TriggeredCompactionHandoff{}, err - } - - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return TriggeredCompactionHandoff{}, fmt.Errorf("begin triggered compaction handoff: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - failure.ProjectID, - failure.AgentID, - failure.RuntimeLockID, - ); err != nil { - return TriggeredCompactionHandoff{}, err - } - - parent, err := loadModelCallContextByID( - ctx, - q, - failure.ProjectID, - failure.AgentID, - input.ParentContextID, - ) - if err != nil { - return TriggeredCompactionHandoff{}, fmt.Errorf("load overflowing model context: %w", err) - } - if parent.OperationKind != ModelCallOperationNormal || parent.State != ModelCallContextStarted || - parent.RuntimeLockID != failure.RuntimeLockID || - input.SourceEventSequenceEnd > parent.InputEventSequence { - return TriggeredCompactionHandoff{}, storeerr.ErrStateTransitionConflict - } - if failure.ModelCallContextID != parent.ID { + f := input.Failure + if input.ParentContextID != f.ModelCallContextID || f.RecoveryKind != ModelCallRecoveryCompact { return TriggeredCompactionHandoff{}, storeerr.ErrStateTransitionConflict } - - summarizedThrough := int64(0) - checkpoint, checkpointErr := q.GetLatestApplicableContextCheckpoint( - ctx, - dbsqlc.GetLatestApplicableContextCheckpointParams{ - ProjectID: failure.ProjectID, - AgentID: failure.AgentID, - MaxEventSequence: parent.InputEventSequence, - }, - ) - if checkpointErr == nil { - summarizedThrough = checkpoint.SummarizedThroughEventSequence - } else if !errors.Is(checkpointErr, pgx.ErrNoRows) { - return TriggeredCompactionHandoff{}, fmt.Errorf( - "load prior checkpoint for triggered compaction: %w", - checkpointErr, - ) - } - if input.SourceEventSequenceEnd <= summarizedThrough { - return TriggeredCompactionHandoff{}, storeerr.ErrStateTransitionConflict - } - - parentContext, err := finishModelCallContextTx(ctx, q, finishModelCallContextInput{ - ProjectID: failure.ProjectID, - AgentID: failure.AgentID, - ModelCallContextID: failure.ModelCallContextID, - RuntimeLockID: failure.RuntimeLockID, - ToState: ModelCallContextFailed, - RecoveryKind: ModelCallRecoveryCompact, - APIFormat: failure.APIFormat, - APIVariant: failure.APIVariant, - ProviderRequestID: failure.ProviderRequestID, - ProviderResponseID: failure.ProviderResponseID, - ErrorKind: failure.ErrorKind, - ErrorCode: failure.ErrorCode, - ErrorMessage: failure.ErrorMessage, - ErrorDetails: failure.ErrorDetails, - Usage: failure.Usage, - ProviderReportedCostUSD: failure.ProviderReportedCostUSD, - ProviderMetadata: failure.ProviderMetadata, - }) + unit, h, err := beginExecution(ctx, s, f.ProjectID, f.AgentID) if err != nil { return TriggeredCompactionHandoff{}, err } - boundaryPreempted, err := supersedeUnacceptedAtLaterModelReadyBoundaryTx( + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.BeginCompaction( ctx, - txNotifications, - tx, - q, - failure.ProjectID, - failure.AgentID, - parent.ID, - "triggered_compaction_preempted", + agentexecution.BeginCompactionInput{ + Failure: executionFailure(f), + SourceEnd: input.SourceEventSequenceEnd, + }, ) if err != nil { return TriggeredCompactionHandoff{}, err } - if boundaryPreempted { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "preempted triggered compaction handoff", - ); err != nil { - return TriggeredCompactionHandoff{}, err - } - return TriggeredCompactionHandoff{ - ParentContext: parentContext, - BoundaryPreempted: true, - }, nil - } - compactionClaimTx, err := claimCompactionContextTx(ctx, q, ClaimCompactionModelCallInput{ - ProjectID: failure.ProjectID, - AgentID: failure.AgentID, - RuntimeLockID: failure.RuntimeLockID, - InputEventSequence: parent.InputEventSequence, - SourceEventSequenceEnd: input.SourceEventSequenceEnd, - ParentContextID: parent.ID, - }) - if err != nil { + if err := applyAdmissionDestination(ctx, + unit, + input.Failure.ProjectID, + input.Failure.AgentID, + result.Admission); err != nil { return TriggeredCompactionHandoff{}, err } - if !compactionClaimTx.created { - return TriggeredCompactionHandoff{}, fmt.Errorf( - "triggered compaction context already exists: %w", - storeerr.ErrStateTransitionConflict, - ) - } - compactionClaim, err := applyModelCallAdmissionTx( + parent, err := loadModelCallContextByID( ctx, - txNotifications, - tx, - q, - compactionClaimTx, - failure.RuntimeLockID, + dbsqlc.New(unit.DB()), + f.ProjectID, + f.AgentID, + input.ParentContextID, ) if err != nil { return TriggeredCompactionHandoff{}, err } - if compactionClaim.Claimed { - if err := q.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: failure.ProjectID, - AgentID: failure.AgentID, - Metadata: json.RawMessage(`{"reason":"triggered_compaction"}`), - }); err != nil { - return TriggeredCompactionHandoff{}, fmt.Errorf( - "reconcile wakeup after triggered compaction: %w", - err, - ) - } + claim, err := executionClaim(ctx, unit, f.ProjectID, f.AgentID, result.Model) + if err != nil { + return TriggeredCompactionHandoff{}, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "triggered compaction handoff", - ); err != nil { + if err = unit.Commit(ctx, "begin compaction"); err != nil { return TriggeredCompactionHandoff{}, err } return TriggeredCompactionHandoff{ - ParentContext: parentContext, - CompactionCall: compactionClaim, + ParentContext: parent, + CompactionCall: claim, + BoundaryPreempted: result.Preempted, }, nil } @@ -210,163 +67,42 @@ func (s *Store) ReplaceCompactionSource( ctx context.Context, input ReplaceCompactionSourceInput, ) (ReplaceCompactionSourceResult, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.ModelCallContextID == uuid.Nil || - input.ErrorKind == "" || input.ErrorMessage == "" || input.NextSourceEventSequenceEnd <= 0 { - return ReplaceCompactionSourceResult{}, errors.New( - "project, agent, runtime, compaction context, error, and replacement source end are required", - ) - } - errorDetails, err := normalizedJSONObject(input.ErrorDetails, "compaction replacement error details") - if err != nil { - return ReplaceCompactionSourceResult{}, err - } - input.ErrorDetails = errorDetails - input.Usage = modelUsageForStorage(input.Usage) - if err := validateModelCallFailureEvidence( - input.APIFormat, - input.APIVariant, - "", - input.ProviderRequestID, - input.ProviderResponseID, - input.Usage != (modelenvelope.Usage{}), - input.ProviderReportedCostUSD, - ); err != nil { - return ReplaceCompactionSourceResult{}, err - } - - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return ReplaceCompactionSourceResult{}, fmt.Errorf("begin compaction source replacement: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return ReplaceCompactionSourceResult{}, err } - contextRow, err := loadModelCallContextByID( + defer func() { _ = unit.Rollback(ctx) }() + failure := agentexecution.ModelFailure{ + ContextID: input.ModelCallContextID, + RuntimeLockID: input.RuntimeLockID, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, APIVariant: input.APIVariant, + RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, Usage: input.Usage, + Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}, + } + result, err := h.BeginCompaction( ctx, - q, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, + agentexecution.BeginCompactionInput{Failure: failure, SourceEnd: input.NextSourceEventSequenceEnd}, ) - if err != nil { - return ReplaceCompactionSourceResult{}, fmt.Errorf("load compaction context for smaller source: %w", err) - } - if contextRow.OperationKind != ModelCallOperationCompaction || - contextRow.State != ModelCallContextStarted || - contextRow.RuntimeLockID != input.RuntimeLockID || - contextRow.SourceEventSequenceEnd == nil || - input.NextSourceEventSequenceEnd >= *contextRow.SourceEventSequenceEnd { - return ReplaceCompactionSourceResult{}, storeerr.ErrStateTransitionConflict - } - sourceStart, err := compactionSourceStartTx(ctx, q, contextRow) if err != nil { return ReplaceCompactionSourceResult{}, err } - if input.NextSourceEventSequenceEnd < sourceStart { - return ReplaceCompactionSourceResult{}, storeerr.ErrStateTransitionConflict - } - if _, err := finishModelCallContextTx(ctx, q, finishModelCallContextInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ToState: ModelCallContextFailed, - RecoveryKind: ModelCallRecoveryReduceCompactionSource, - APIFormat: input.APIFormat, - APIVariant: input.APIVariant, - ProviderRequestID: input.ProviderRequestID, - ProviderResponseID: input.ProviderResponseID, - ErrorKind: input.ErrorKind, - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: input.ErrorDetails, - Usage: input.Usage, - ProviderReportedCostUSD: input.ProviderReportedCostUSD, - ProviderMetadata: input.ProviderMetadata, - }); err != nil { - return ReplaceCompactionSourceResult{}, err - } - boundaryPreempted, err := supersedeUnacceptedAtLaterModelReadyBoundaryTx( - ctx, - txNotifications, - tx, - q, + if err := applyAdmissionDestination(ctx, + unit, input.ProjectID, input.AgentID, - contextRow.ID, - "compaction_source_preempted", - ) - if err != nil { + result.Admission); err != nil { return ReplaceCompactionSourceResult{}, err } - if boundaryPreempted { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "preempted compaction source replacement", - ); err != nil { - return ReplaceCompactionSourceResult{}, err - } - return ReplaceCompactionSourceResult{BoundaryPreempted: true}, nil - } - - parentID, err := q.GetNormalModelCallContextByIdentity( - ctx, - dbsqlc.GetNormalModelCallContextByIdentityParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - InputEventSequence: contextRow.InputEventSequence, - }, - ) - if err != nil { - return ReplaceCompactionSourceResult{}, fmt.Errorf("load parent normal context for replacement compaction: %w", err) - } - nextClaimTx, err := claimCompactionContextTx(ctx, q, ClaimCompactionModelCallInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - RuntimeLockID: input.RuntimeLockID, - InputEventSequence: contextRow.InputEventSequence, - SourceEventSequenceEnd: input.NextSourceEventSequenceEnd, - ParentContextID: parentID, - }) + claim, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, result.Model) if err != nil { return ReplaceCompactionSourceResult{}, err } - if !nextClaimTx.created { - return ReplaceCompactionSourceResult{}, fmt.Errorf( - "replacement compaction context already exists: %w", - storeerr.ErrStateTransitionConflict, - ) - } - nextClaim, err := applyModelCallAdmissionTx( - ctx, - txNotifications, - tx, - q, - nextClaimTx, - input.RuntimeLockID, - ) - if err != nil { - return ReplaceCompactionSourceResult{}, err - } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "compaction source replacement", - ); err != nil { + if err = unit.Commit(ctx, "reduce compaction source"); err != nil { return ReplaceCompactionSourceResult{}, err } - return ReplaceCompactionSourceResult{CompactionCall: nextClaim}, nil + return ReplaceCompactionSourceResult{CompactionCall: claim, BoundaryPreempted: result.Preempted}, nil } diff --git a/internal/storage/executionstore/model_contexts_lifecycle.go b/internal/storage/executionstore/model_contexts_lifecycle.go index c6252c081..51ff023c1 100644 --- a/internal/storage/executionstore/model_contexts_lifecycle.go +++ b/internal/storage/executionstore/model_contexts_lifecycle.go @@ -2,18 +2,9 @@ package executionstore import ( "context" - "encoding/json" "errors" - "fmt" - "time" - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" ) func (s *Store) RecordRetryableModelCallFailure( @@ -24,251 +15,28 @@ func (s *Store) RecordRetryableModelCallFailure( input.RecoveryKind = ModelCallRecoveryRetry } if input.RecoveryKind != ModelCallRecoveryRetry { - return ModelCallContextRecord{}, fmt.Errorf( - "retryable model call failure requires recovery kind %q", - ModelCallRecoveryRetry, - ) + return ModelCallContextRecord{}, errors.New("retry recovery is required") } - if err := validateRecoverableModelCallFailure(input); err != nil { - return ModelCallContextRecord{}, err - } - var err error - input.ErrorDetails, err = normalizedJSONObject(input.ErrorDetails, "model call error details") + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { return ModelCallContextRecord{}, err } - input.Usage = modelUsageForStorage(input.Usage) - if err := validateModelCallFailureEvidence( - input.APIFormat, - input.APIVariant, - "", - input.ProviderRequestID, - input.ProviderResponseID, - input.Usage != (modelenvelope.Usage{}), - input.ProviderReportedCostUSD, - ); err != nil { - return ModelCallContextRecord{}, err - } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return ModelCallContextRecord{}, fmt.Errorf("begin retryable model call failure: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { + defer func() { _ = unit.Rollback(ctx) }() + if _, err = h.FailModel(ctx, executionFailure(input)); err != nil { return ModelCallContextRecord{}, err } - currentContext, err := loadModelCallContextByID( + record, err := loadModelCallContextByID( ctx, - q, + dbsqlc.New(unit.DB()), input.ProjectID, input.AgentID, input.ModelCallContextID, ) - if err != nil { - return ModelCallContextRecord{}, fmt.Errorf("load retryable model call context: %w", err) - } - if currentContext.State != ModelCallContextStarted || - currentContext.AttemptNumber > MaxModelCallRetriesPerOperation { - return ModelCallContextRecord{}, storeerr.ErrStateTransitionConflict - } - contextRecord, err := finishModelCallContextTx(ctx, q, finishModelCallContextInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ToState: ModelCallContextFailed, - RecoveryKind: ModelCallRecoveryRetry, - APIFormat: input.APIFormat, - APIVariant: input.APIVariant, - ProviderRequestID: input.ProviderRequestID, - ProviderResponseID: input.ProviderResponseID, - ErrorKind: input.ErrorKind, - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: input.ErrorDetails, - RetryDelay: &input.RetryDelay, - Usage: input.Usage, - ProviderReportedCostUSD: input.ProviderReportedCostUSD, - ProviderMetadata: input.ProviderMetadata, - }) if err != nil { return ModelCallContextRecord{}, err } - if err := q.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - Metadata: json.RawMessage(`{"reason":"model_call_retry"}`), - ProjectID: input.ProjectID, - AgentID: input.AgentID, - }); err != nil { - return ModelCallContextRecord{}, fmt.Errorf("reconcile model call retry wakeup: %w", err) - } - if err := tx.Commit(ctx); err != nil { - return ModelCallContextRecord{}, fmt.Errorf("commit retryable model call failure: %w", err) - } - return contextRecord, nil -} - -func validateRecoverableModelCallFailure(input RecordRecoverableModelCallFailureInput) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ModelCallContextID == uuid.Nil || - input.RuntimeLockID == uuid.Nil || input.ErrorKind == "" || - input.ErrorMessage == "" { - return errors.New("project, agent, context, runtime, and error are required") - } - if input.RetryDelay < 0 { - return errors.New("model call retry delay cannot be negative") - } - if input.RecoveryKind != ModelCallRecoveryRetry && input.RecoveryKind != ModelCallRecoveryCompact { - return fmt.Errorf("unsupported model call recovery kind %q", input.RecoveryKind) - } - return nil -} - -func validateModelCallFailureEvidence( - apiFormat modelprotocol.APIFormat, - apiVariant modelprotocol.APIVariant, - servedProviderModelSlug, - providerRequestID, - providerResponseID string, - hasUsage bool, - providerReportedCostUSD modelenvelope.ProviderReportedCostUSD, -) error { - if err := modelenvelope.ValidateProviderReportedCostUSD(providerReportedCostUSD); err != nil { - return fmt.Errorf("provider-reported cost: %w", err) - } - if (apiFormat == "") != (apiVariant == "") { - return errors.New("model call API format and variant must be recorded together") - } - if apiFormat == "" && (servedProviderModelSlug != "" || providerRequestID != "" || - providerResponseID != "" || hasUsage || providerReportedCostUSD != "") { - return errors.New("provider evidence requires a model call API identity") - } - return nil -} - -type finishModelCallContextInput struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelCallContextID uuid.UUID - RuntimeLockID uuid.UUID - ToState ModelCallState - RecoveryKind ModelCallRecoveryKind - APIFormat modelprotocol.APIFormat - APIVariant modelprotocol.APIVariant - ProviderRequestID string - ProviderResponseID string - ErrorKind modelprotocol.ErrorKind - ErrorCode string - ErrorMessage string - ErrorDetails json.RawMessage - RetryDelay *time.Duration - Usage modelenvelope.Usage - ProviderReportedCostUSD modelenvelope.ProviderReportedCostUSD - ProviderMetadata modelenvelope.ProviderMetadata -} - -type modelCallContextRuntimeAuthority uint8 - -const ( - modelCallContextRuntimeOwned modelCallContextRuntimeAuthority = iota - modelCallContextRuntimeTeardown -) - -func (a modelCallContextRuntimeAuthority) allowsInactiveRuntimeLockForTeardown() (bool, error) { - switch a { - case modelCallContextRuntimeOwned: - return false, nil - case modelCallContextRuntimeTeardown: - return true, nil - default: - return false, fmt.Errorf("unsupported model call context runtime authority %d", a) - } -} - -func finishModelCallContextWithAuthorityTx( - ctx context.Context, - q *dbsqlc.Queries, - input finishModelCallContextInput, - authority modelCallContextRuntimeAuthority, -) (ModelCallContextRecord, error) { - allowInactiveRuntimeLockForTeardown, err := authority.allowsInactiveRuntimeLockForTeardown() - if err != nil { + if err = unit.Commit(ctx, "record retryable failure"); err != nil { return ModelCallContextRecord{}, err } - usage := usageColumnsFromModelUsage(input.Usage) - if err := modelenvelope.ValidateProviderReportedCostUSD(input.ProviderReportedCostUSD); err != nil { - return ModelCallContextRecord{}, fmt.Errorf("provider-reported cost: %w", err) - } - providerMetadata, err := json.Marshal(input.ProviderMetadata) - if err != nil { - return ModelCallContextRecord{}, fmt.Errorf("encode provider metadata: %w", err) - } - var retryDelayMicroseconds *int64 - if input.RetryDelay != nil { - if *input.RetryDelay < 0 { - return ModelCallContextRecord{}, errors.New("model call retry delay cannot be negative") - } - delay := input.RetryDelay.Microseconds() - retryDelayMicroseconds = &delay - } - id, err := q.FinishModelCallContext(ctx, dbsqlc.FinishModelCallContextParams{ - ToState: string(input.ToState), - RecoveryKind: storeutil.TextFromEmpty(string(input.RecoveryKind)), - ApiFormat: string(input.APIFormat), - ApiVariant: string(input.APIVariant), - ProviderRequestID: input.ProviderRequestID, - ProviderResponseID: input.ProviderResponseID, - ErrorKind: string(input.ErrorKind), - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: normalizedJSONOrObject(input.ErrorDetails), - RetryDelayMicroseconds: retryDelayMicroseconds, - InputTokensTotal: storeutil.Int32Ptr(usage.InputTokens), - UncachedInputTokens: storeutil.Int32Ptr(usage.UncachedInputTokens), - CacheReadInputTokens: storeutil.Int32Ptr(usage.CacheReadTokens), - CacheWriteInputTokens: storeutil.Int32Ptr(usage.CacheWriteTokens), - OutputTokensTotal: storeutil.Int32Ptr(usage.OutputTokens), - ReasoningOutputTokens: storeutil.Int32Ptr(usage.ReasoningTokens), - ProviderReportedCostUsd: providerReportedCostUSDToSQLC(input.ProviderReportedCostUSD), - ProviderMetadata: providerMetadata, - ID: input.ModelCallContextID, - ProjectID: input.ProjectID, - AgentID: input.AgentID, - RuntimeLockID: input.RuntimeLockID, - AllowInactiveRuntimeLockForTeardown: allowInactiveRuntimeLockForTeardown, - }) - if errors.Is(err, pgx.ErrNoRows) { - if allowInactiveRuntimeLockForTeardown { - return ModelCallContextRecord{}, storeerr.ErrStateTransitionConflict - } - return ModelCallContextRecord{}, storeerr.ErrRuntimeLockInactive - } - if err != nil { - return ModelCallContextRecord{}, fmt.Errorf("finish model call context: %w", err) - } - record, err := loadModelCallContextByID(ctx, q, input.ProjectID, input.AgentID, id) - if err != nil { - return ModelCallContextRecord{}, fmt.Errorf("load finished model call context: %w", err) - } return record, nil } - -func finishModelCallContextTx( - ctx context.Context, - q *dbsqlc.Queries, - input finishModelCallContextInput, -) (ModelCallContextRecord, error) { - return finishModelCallContextWithAuthorityTx(ctx, q, input, modelCallContextRuntimeOwned) -} - -func normalizedJSONOrObject(value json.RawMessage) json.RawMessage { - if len(value) == 0 || string(value) == "null" { - return json.RawMessage(`{}`) - } - return normalizedJSON(value) -} diff --git a/internal/storage/executionstore/model_contexts_lifecycle_test.go b/internal/storage/executionstore/model_contexts_lifecycle_test.go deleted file mode 100644 index 84c3a1eb6..000000000 --- a/internal/storage/executionstore/model_contexts_lifecycle_test.go +++ /dev/null @@ -1,80 +0,0 @@ -package executionstore - -import ( - "testing" - - "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/modelprotocol" -) - -func TestValidateModelCallFailureEvidence(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - servedProviderModelSlug string - providerRequestID string - providerResponseID string - hasUsage bool - providerReportedCostUSD modelenvelope.ProviderReportedCostUSD - }{ - {name: "served model", servedProviderModelSlug: "model"}, - {name: "request ID", providerRequestID: "req"}, - {name: "response ID", providerResponseID: "resp"}, - {name: "usage", hasUsage: true}, - {name: "provider-reported cost", providerReportedCostUSD: "0.01"}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - t.Parallel() - err := validateModelCallFailureEvidence( - "", - "", - test.servedProviderModelSlug, - test.providerRequestID, - test.providerResponseID, - test.hasUsage, - test.providerReportedCostUSD, - ) - if err == nil { - t.Fatal("provider evidence without API identity was accepted") - } - }) - } - - if err := validateModelCallFailureEvidence( - modelprotocol.APIFormatOpenAIResponses, - modelprotocol.APIVariantDefault, - "model", - "req", - "resp", - true, - "", - ); err != nil { - t.Fatalf("provider evidence with API identity was rejected: %v", err) - } - - if err := validateModelCallFailureEvidence( - modelprotocol.APIFormatOpenAIResponses, - "", - "", - "", - "", - false, - "", - ); err == nil { - t.Fatal("partial API identity was accepted") - } - - if err := validateModelCallFailureEvidence( - modelprotocol.APIFormatOpenAIResponses, - modelprotocol.APIVariantDefault, - "", - "", - "", - false, - "not-a-decimal", - ); err == nil { - t.Fatal("invalid provider-reported cost was accepted") - } -} diff --git a/internal/storage/executionstore/model_contexts_preparation.go b/internal/storage/executionstore/model_contexts_preparation.go new file mode 100644 index 000000000..4d2fb906c --- /dev/null +++ b/internal/storage/executionstore/model_contexts_preparation.go @@ -0,0 +1,45 @@ +package executionstore + +import ( + "context" + "errors" + + "github.com/google/uuid" +) + +type PrepareNormalModelCallInput struct { + ProjectID uuid.UUID + AgentID uuid.UUID + RuntimeLockID uuid.UUID + OpeningInputIDs []uuid.UUID + SourceModelCallContextID uuid.UUID + SourceModelOutputID uuid.UUID +} + +type PreparedNormalModelCall struct { + Claim ModelCallClaim + Snapshot AgentConfigSnapshotRecord + ContextData *ModelContextData +} + +func (s *Store) PrepareNormalModelCall( + ctx context.Context, + input PrepareNormalModelCallInput, +) (PreparedNormalModelCall, error) { + if input.RuntimeLockID == uuid.Nil || len(input.OpeningInputIDs) == 0 { + return PreparedNormalModelCall{}, errors.New("runtime and opening inputs are required") + } + unit, _, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return PreparedNormalModelCall{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := prepareExecutionModel(ctx, unit, input) + if err != nil { + return PreparedNormalModelCall{}, err + } + if err = unit.Commit(ctx, "prepare model call"); err != nil { + return PreparedNormalModelCall{}, err + } + return result, nil +} diff --git a/internal/storage/executionstore/model_contexts_preparation_integration_test.go b/internal/storage/executionstore/model_contexts_preparation_integration_test.go new file mode 100644 index 000000000..c99b816e2 --- /dev/null +++ b/internal/storage/executionstore/model_contexts_preparation_integration_test.go @@ -0,0 +1,140 @@ +//go:build integration + +package executionstore_test + +import ( + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + "github.com/stretchr/testify/require" +) + +func TestPrepareNormalModelCallUsesSnapshotAfterLockWait(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "prepare_snapshot") + input := ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "initial") + blocker, err := f.Store.Execution().IntegrationBeginUnit(ctx) + require.NoError(t, err) + defer func() { _ = blocker.Rollback(ctx) }() + _, err = blocker.LockAgent( + ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: f.AgentID}, + agentexecution.ExternalAuthority{}, + ) + require.NoError(t, err) + var pid int32 + require.NoError(t, blocker.DB().QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&pid)) + type result struct { + prepared executionstore.PreparedNormalModelCall + err error + } + done := make(chan result, 1) + go func() { + prepared, err := f.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: f.AgentID, + RuntimeLockID: f.Lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, + }) + done <- result{prepared, err} + }() + integrationdb.WaitForLockWaitBlockedBy(t, ctx, f.Store.pool, "", pid) + admitted, err := executionstore.IntegrationAdmitInputs(ctx, blocker, testProjectID, f.AgentID) + require.NoError(t, err) + require.Len(t, admitted.Events, 1) + sequence := admitted.Events[0].Sequence + require.NoError(t, blocker.Commit(ctx, "admit after wait")) + got := <-done + require.NoError(t, got.err) + require.True(t, got.prepared.Claim.Claimed) + require.Equal(t, sequence, got.prepared.Snapshot.InputEventSequence) + require.Equal(t, sequence, got.prepared.Claim.Context.InputEventSequence) + require.Equal(t, got.prepared.Snapshot.AgentConfig.ID, got.prepared.Claim.Context.AgentConfigID) +} + +func TestPrepareNormalModelCallRollbackAndReplay(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "prepare_rollback") + input := ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "initial") + _, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + prepare := executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: f.AgentID, + RuntimeLockID: f.Lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, + } + _, err = f.Store.pool.Exec( + ctx, + `CREATE FUNCTION reject_prepared_context() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN RAISE EXCEPTION 'reject'; END $$; +CREATE CONSTRAINT TRIGGER reject_prepared_context AFTER INSERT ON model_call_contexts +DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION reject_prepared_context();`, + ) + require.NoError(t, err) + _, err = f.Store.Execution().PrepareNormalModelCall(ctx, prepare) + require.ErrorContains(t, err, "commit prepare model call") + var count int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM model_call_contexts WHERE agent_id=$1`, f.AgentID, + ).Scan(&count)) + require.Zero(t, count) + _, err = f.Store.pool.Exec(ctx, `DROP TRIGGER reject_prepared_context ON model_call_contexts`) + require.NoError(t, err) + prepared, err := f.Store.Execution().PrepareNormalModelCall(ctx, prepare) + require.NoError(t, err) + require.True(t, prepared.Claim.Claimed) + require.Equal(t, prepared.Snapshot.AgentConfig.ID, prepared.Claim.Context.AgentConfigID) + require.Equal(t, prepared.Snapshot.InputEventSequence, prepared.Claim.Context.InputEventSequence) + replay, err := f.Store.Execution().PrepareNormalModelCall(ctx, prepare) + require.NoError(t, err) + require.False(t, replay.Claim.Created) + require.False(t, replay.Claim.Claimed) + require.Equal(t, prepared.Claim.Context.ID, replay.Claim.Context.ID) +} + +func TestPrepareNormalModelCallValidatesOwnershipAndSelectedWork(t *testing.T) { + for _, mode := range []string{"expired", "opening inputs", "missing source", "wrong source"} { + t.Run(mode, func(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "prepare_validation") + id := createToolCallForProcessTest(t, ctx, f, "tool", "read_process") + ownedCompleteTool(t, ctx, f, id) + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + input := executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, + OpeningInputIDs: work.Model.InputIDs, + SourceModelCallContextID: work.Model.SourceModelCallContextID, + SourceModelOutputID: work.Model.SourceModelOutputID, + } + wantErr := storeerr.ErrAgentNotAdvanceable + switch mode { + case "expired": + expireAgentRuntimeLockForTest(t, ctx, f.Store, f.Lock.ID) + wantErr = storeerr.ErrRuntimeLockInactive + case "opening inputs": + input.OpeningInputIDs = []uuid.UUID{uuid.New()} + case "missing source": + input.SourceModelCallContextID, input.SourceModelOutputID = uuid.Nil, uuid.Nil + case "wrong source": + input.SourceModelOutputID = uuid.New() + } + _, err = f.Store.Execution().PrepareNormalModelCall(ctx, input) + require.ErrorIs(t, err, wantErr) + var count int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM model_call_contexts WHERE agent_id=$1 AND state='started'`, f.AgentID, + ).Scan(&count)) + require.Zero(t, count) + }) + } +} diff --git a/internal/storage/executionstore/model_contexts_replay_integration_test.go b/internal/storage/executionstore/model_contexts_replay_integration_test.go index 807440b64..7486cef2a 100644 --- a/internal/storage/executionstore/model_contexts_replay_integration_test.go +++ b/internal/storage/executionstore/model_contexts_replay_integration_test.go @@ -38,9 +38,13 @@ func TestProviderReplaySuppressionCutoffUsesCompatiblePriorFailures(t *testing.T INSERT INTO model_call_contexts( org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, - runtime_lock_id, state, created_at + runtime_lock_id, state, created_at,turn_id,opening_input_ids,opening_event_sequence ) -VALUES ($1, $2, $3, 'normal', 1, $4, $5, $6, $7, 'started', statement_timestamp()) +VALUES ($1, $2, $3, 'normal', 1, $4, $5, $6, $7, 'started', statement_timestamp(), +(SELECT current_turn_id FROM agent_execution_state WHERE agent_id=$3), (SELECT +array_agg(agent_input_id ORDER BY sequence) FROM agent_events WHERE agent_id=$3 AND +is_opening_event AND event_kind='agent_input'), (SELECT min(sequence) FROM agent_events +WHERE agent_id=$3 AND is_opening_event)) RETURNING id `, testOrgID, testProjectID, fixture.AgentID, agent.CurrentConfigID, revisionID, frontier, fixture.Lock.ID).Scan(&contextID); err != nil { @@ -55,9 +59,13 @@ RETURNING id INSERT INTO model_call_contexts( org_id, project_id, agent_id, operation_kind, attempt_number, agent_config_id, configured_model_revision_id, input_event_sequence, - source_event_sequence_end, runtime_lock_id, state, created_at + source_event_sequence_end, runtime_lock_id, state, created_at,turn_id,opening_input_ids,opening_event_sequence ) -VALUES ($1, $2, $3, 'compaction', 1, $4, $5, $6, 1, $7, 'started', statement_timestamp()) +VALUES ($1, $2, $3, 'compaction', 1, $4, $5, $6, 1, $7, 'started', statement_timestamp(), +(SELECT current_turn_id FROM agent_execution_state WHERE agent_id=$3), (SELECT +array_agg(agent_input_id ORDER BY sequence) FROM agent_events WHERE agent_id=$3 AND +is_opening_event AND event_kind='agent_input'), (SELECT min(sequence) FROM agent_events +WHERE agent_id=$3 AND is_opening_event)) RETURNING id `, testOrgID, testProjectID, fixture.AgentID, agent.CurrentConfigID, revisionID, frontier, fixture.Lock.ID).Scan(&contextID); err != nil { diff --git a/internal/storage/executionstore/model_contexts_store.go b/internal/storage/executionstore/model_contexts_store.go index f3e36d31c..c23c4d7e7 100644 --- a/internal/storage/executionstore/model_contexts_store.go +++ b/internal/storage/executionstore/model_contexts_store.go @@ -2,8 +2,6 @@ package executionstore import ( "encoding/json" - "fmt" - "hash/fnv" "time" "github.com/google/uuid" @@ -14,9 +12,6 @@ import ( const ( MaxModelCallRetriesPerOperation = 8 - baseModelCallRetryBackoff = time.Second - maxModelCallRetryBackoff = 30 * time.Second - ModelCallOperationNormal ModelCallOperation = "normal" ModelCallOperationCompaction ModelCallOperation = "compaction" @@ -66,46 +61,12 @@ type ModelCallContextRecord struct { CompletedAt *time.Time `json:"completed_at,omitempty"` } -func ModelCallRetryBackoff(attemptNumber int, contextID string) time.Duration { - if attemptNumber < 1 { - attemptNumber = 1 - } - delay := baseModelCallRetryBackoff - for i := 1; i < attemptNumber && delay < maxModelCallRetryBackoff; i++ { - delay *= 2 - } - - percent := deterministicModelCallRetryJitterPercent(contextID, attemptNumber) - delay = time.Duration(int64(delay) * int64(percent) / 100) - if delay > maxModelCallRetryBackoff { - return maxModelCallRetryBackoff - } - return delay -} - -func deterministicModelCallRetryJitterPercent(contextID string, attemptNumber int) int { - hash := fnv.New32a() - _, _ = hash.Write([]byte(fmt.Sprintf("%s:%d", contextID, attemptNumber))) - return 80 + int(hash.Sum32()%41) -} - type ModelCallClaim struct { Context ModelCallContextRecord Created bool Claimed bool } -type ClaimNormalModelCallInput struct { - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID - OpeningInputIDs []uuid.UUID - AgentConfigID uuid.UUID - InputEventSequence int64 - SourceModelCallContextID uuid.UUID - SourceModelOutputID uuid.UUID -} - type ClaimCompactionModelCallInput struct { ProjectID uuid.UUID AgentID uuid.UUID diff --git a/internal/storage/executionstore/model_outputs_content_projection.go b/internal/storage/executionstore/model_outputs_content_projection.go deleted file mode 100644 index 1cd82e02e..000000000 --- a/internal/storage/executionstore/model_outputs_content_projection.go +++ /dev/null @@ -1,219 +0,0 @@ -package executionstore - -import ( - "context" - "fmt" - - "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" -) - -type toolCallContentBlockArgs struct { - ToolCallID uuid.UUID -} - -func createModelOutputContentBlocksTx( - ctx context.Context, - tx pgx.Tx, - contextRow ModelCallContextRecord, - envelope modelenvelope.ResponseEnvelope, - modelOutputID uuid.UUID, - toolCalls map[string]toolCallContentBlockArgs, -) (map[string]uuid.UUID, error) { - contentBlockByProviderCallID := make(map[string]uuid.UUID, len(toolCalls)) - if len(envelope.Normalized.Content) == 0 { - return contentBlockByProviderCallID, nil - } - for index, part := range envelope.Normalized.Content { - ordinal := int32(index) - switch part.Type { - case modelenvelope.ResponsePartTypeText, - modelenvelope.ResponsePartTypeReasoning, - modelenvelope.ResponsePartTypeError: - if part.Text == "" { - continue - } - blockKind := ContentBlockKindText - if part.Type == modelenvelope.ResponsePartTypeReasoning { - blockKind = ContentBlockKindReasoning - } else if part.Type == modelenvelope.ResponsePartTypeError { - blockKind = ContentBlockKindError - } - if _, err := createContentBlockTx(ctx, tx, CreateContentBlockInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - OwnerKind: ContentBlockOwnerModelOutput, - OwnerModelOutputID: modelOutputID, - Ordinal: ordinal, - BlockKind: blockKind, - TextContent: part.Text, - }); err != nil { - return nil, err - } - case modelenvelope.ResponsePartTypeToolCall: - args, ok := toolCalls[part.ProviderCallID] - if !ok || args.ToolCallID == uuid.Nil { - return nil, fmt.Errorf( - "model output content tool_call part %d (provider_call_id=%q) has no recorded tool call", - index, - part.ProviderCallID, - ) - } - block, err := createContentBlockTx(ctx, tx, CreateContentBlockInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - OwnerKind: ContentBlockOwnerModelOutput, - OwnerModelOutputID: modelOutputID, - Ordinal: ordinal, - BlockKind: ContentBlockKindToolCall, - ToolCallID: args.ToolCallID, - }) - if err != nil { - return nil, err - } - contentBlockByProviderCallID[part.ProviderCallID] = block.ID - default: - return nil, fmt.Errorf("model output content part %d has unsupported type %q", index, part.Type) - } - } - return contentBlockByProviderCallID, nil -} - -func modelOutputHasContentBlocksTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID, modelOutputID uuid.UUID, -) (bool, error) { - rows, err := dbsqlc.New(tx). - ListContentBlocksForModelOutput(ctx, dbsqlc.ListContentBlocksForModelOutputParams{ - ProjectID: projectID, - AgentID: agentID, - ModelOutputID: storeutil.IDFromNil(modelOutputID), - }) - if err != nil { - return false, fmt.Errorf("list model output content blocks: %w", err) - } - return len(rows) > 0, nil -} - -func validateModelOutputContentReplayTx( - ctx context.Context, - tx pgx.Tx, - projectID, agentID, modelOutputID uuid.UUID, - envelope modelenvelope.ResponseEnvelope, -) error { - rows, err := dbsqlc.New(tx). - ListContentBlocksForModelOutput(ctx, dbsqlc.ListContentBlocksForModelOutputParams{ - ProjectID: projectID, - AgentID: agentID, - ModelOutputID: storeutil.IDFromNil(modelOutputID), - }) - if err != nil { - return fmt.Errorf("list model output content blocks for replay: %w", err) - } - expected, err := modelOutputReplayParts(envelope) - if err != nil { - return err - } - if len(rows) != len(expected) { - return storeerr.ErrIdempotencyConflict - } - for index, part := range expected { - row := rows[index] - providerCallID := "" - if part.Kind == ContentBlockKindToolCall { - providerCallID, err = dbsqlc.New(tx).GetToolCallProviderCallID( - ctx, - dbsqlc.GetToolCallProviderCallIDParams{ - ProjectID: projectID, - AgentID: agentID, - ID: storeutil.IDFromPtr(row.ToolCallID), - }, - ) - if err != nil { - return fmt.Errorf("load replayed tool call provider id: %w", err) - } - } - if row.Ordinal != part.Ordinal || ContentBlockKind(row.BlockKind) != part.Kind || - row.TextContent != part.Text || storeutil.IDFromPtr(row.ArtifactID) != uuid.Nil || - providerCallID != part.ProviderCallID { - return storeerr.ErrIdempotencyConflict - } - } - return nil -} - -type modelOutputReplayPart struct { - Ordinal int32 - Kind ContentBlockKind - Text string - ProviderCallID string -} - -func modelOutputReplayParts(envelope modelenvelope.ResponseEnvelope) ([]modelOutputReplayPart, error) { - parts := make([]modelOutputReplayPart, 0, len(envelope.Normalized.Content)) - for index, part := range envelope.Normalized.Content { - replayPart := modelOutputReplayPart{ - Ordinal: int32(index), - } - switch part.Type { - case modelenvelope.ResponsePartTypeText, - modelenvelope.ResponsePartTypeReasoning, - modelenvelope.ResponsePartTypeError: - if part.Text == "" { - continue - } - replayPart.Kind = ContentBlockKindText - if part.Type == modelenvelope.ResponsePartTypeReasoning { - replayPart.Kind = ContentBlockKindReasoning - } else if part.Type == modelenvelope.ResponsePartTypeError { - replayPart.Kind = ContentBlockKindError - } - replayPart.Text = part.Text - case modelenvelope.ResponsePartTypeToolCall: - replayPart.Kind = ContentBlockKindToolCall - replayPart.ProviderCallID = part.ProviderCallID - default: - return nil, fmt.Errorf("model output content part %d has unsupported type %q", index, part.Type) - } - parts = append(parts, replayPart) - } - return parts, nil -} - -func modelOutputAuthorityInputFromContext( - input RecordModelOutputAndCompleteContextInput, -) CreateModelOutputAuthorityInput { - servedProviderModelSlug := "" - envelope := input.ProviderResponse - if envelope.ServedProviderModelSlug != "" { - servedProviderModelSlug = envelope.ServedProviderModelSlug - } - return CreateModelOutputAuthorityInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - ServedProviderModelSlug: servedProviderModelSlug, - StopReason: envelope.Normalized.StopReason, - ProviderReplay: envelope.ProviderReplay, - Usage: envelope.Normalized.Usage, - } -} - -func modelOutputAuthorityInputFromToolSourceContext( - input RecordToolCallSourceAndCompleteContextInput, -) CreateModelOutputAuthorityInput { - return modelOutputAuthorityInputFromContext( - RecordModelOutputAndCompleteContextInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - RuntimeLockID: input.RuntimeLockID, - ModelCallContextID: input.ModelCallContextID, - ProviderResponse: input.ProviderResponse, - }, - ) -} diff --git a/internal/storage/executionstore/model_outputs_terminal_error.go b/internal/storage/executionstore/model_outputs_terminal_error.go index b4497d22d..bbb165344 100644 --- a/internal/storage/executionstore/model_outputs_terminal_error.go +++ b/internal/storage/executionstore/model_outputs_terminal_error.go @@ -3,18 +3,13 @@ package executionstore import ( "context" "encoding/json" - "errors" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/modelprotocol" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" ) type RecordModelCallErrorAndCompleteContextInput struct { @@ -40,297 +35,46 @@ func (s *Store) RecordModelCallErrorAndCompleteContext( ctx context.Context, input RecordModelCallErrorAndCompleteContextInput, ) (events.Event, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.ModelCallContextID == uuid.Nil || - input.ErrorKind == "" || input.ErrorMessage == "" { - return events.Event{}, errors.New( - "project, agent, runtime, context, and error are required", - ) - } - var err error - input.ErrorDetails, err = normalizedJSONObject(input.ErrorDetails, "model call error details") - if err != nil { - return events.Event{}, err - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return events.Event{}, fmt.Errorf("begin record model call error: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - q := dbsqlc.New(tx) - if err := lockAgentWithParentTx(ctx, tx, q, input.ProjectID, input.AgentID); err != nil { - return events.Event{}, err - } - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return events.Event{}, err } - result, err := recordTerminalModelCallFailureTx( + defer func() { _ = unit.Rollback(ctx) }() + _, err = h.FailModel( ctx, - txNotifications, - tx, - q, - input, - modelCallContextRuntimeOwned, - ModelCallOperationNormal, + agentexecution.ModelFailure{ContextID: input.ModelCallContextID, RuntimeLockID: input.RuntimeLockID, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + ServedModel: input.ServedProviderModelSlug, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, + APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, + Usage: input.Usage, Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}}, ) if err != nil { return events.Event{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "record model call error"); err != nil { - return events.Event{}, err - } - return result.event, nil -} - -type terminalModelCallFailureResult struct { - context ModelCallContextRecord - event events.Event -} - -func recordTerminalModelCallFailureTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - q *dbsqlc.Queries, - input RecordModelCallErrorAndCompleteContextInput, - runtimeAuthority modelCallContextRuntimeAuthority, - operationKind ModelCallOperation, -) (terminalModelCallFailureResult, error) { - contextRow, err := loadModelCallContextByIDTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) - if err != nil { - return terminalModelCallFailureResult{}, err - } - if contextRow.OperationKind != operationKind { - return terminalModelCallFailureResult{}, storeerr.ErrStateTransitionConflict - } - if contextRow.RuntimeLockID != input.RuntimeLockID { - return terminalModelCallFailureResult{}, storeerr.ErrRuntimeLockInactive - } - normalizedUsage := modelUsageForStorage(input.Usage) - if err := validateModelCallFailureEvidence( - input.APIFormat, - input.APIVariant, - input.ServedProviderModelSlug, - input.ProviderRequestID, - input.ProviderResponseID, - normalizedUsage != (modelenvelope.Usage{}), - input.ProviderReportedCostUSD, - ); err != nil { - return terminalModelCallFailureResult{}, err - } - if contextRow.State == ModelCallContextFailed { - if !sameTerminalModelCallErrorContextIntent(contextRow, input, normalizedUsage) { - return terminalModelCallFailureResult{}, storeerr.ErrIdempotencyConflict - } - event, replayErr := replayModelCallErrorOutputTx(ctx, tx, q, contextRow, input) - if replayErr != nil { - return terminalModelCallFailureResult{}, replayErr - } - return terminalModelCallFailureResult{context: contextRow, event: event.Event}, nil - } - if contextRow.State != ModelCallContextStarted { - return terminalModelCallFailureResult{}, storeerr.ErrRuntimeLockInactive - } - finishedContext, err := finishModelCallContextWithAuthorityTx(ctx, q, finishModelCallContextInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ToState: ModelCallContextFailed, - APIFormat: input.APIFormat, - APIVariant: input.APIVariant, - ProviderRequestID: input.ProviderRequestID, - ProviderResponseID: input.ProviderResponseID, - ErrorKind: input.ErrorKind, - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: input.ErrorDetails, - Usage: normalizedUsage, - ProviderReportedCostUSD: input.ProviderReportedCostUSD, - ProviderMetadata: input.ProviderMetadata, - }, runtimeAuthority) - if err != nil { - return terminalModelCallFailureResult{}, err - } - eventRecord, err := publishModelCallErrorOutputTx( + q := dbsqlc.New(unit.DB()) + output, err := q.GetModelOutputByModelContext( ctx, - txNotifications, - tx, - q, - finishedContext, - modelCallErrorOutputInput{ - ServedProviderModelSlug: input.ServedProviderModelSlug, - ErrorMessage: input.ErrorMessage, - Usage: normalizedUsage, - }, + dbsqlc.GetModelOutputByModelContextParams{ProjectID: input.ProjectID, + AgentID: input.AgentID, ModelCallContextID: input.ModelCallContextID}, ) if err != nil { - return terminalModelCallFailureResult{}, err - } - return terminalModelCallFailureResult{ - context: finishedContext, - event: eventRecord.Event, - }, nil -} - -func sameTerminalModelCallErrorContextIntent( - contextRecord ModelCallContextRecord, - input RecordModelCallErrorAndCompleteContextInput, - usage modelenvelope.Usage, -) bool { - return contextRecord.State == ModelCallContextFailed && - contextRecord.RecoveryKind == "" && - contextRecord.APIFormat == input.APIFormat && - contextRecord.APIVariant == input.APIVariant && - contextRecord.ProviderRequestID == input.ProviderRequestID && - contextRecord.ProviderResponseID == input.ProviderResponseID && - contextRecord.ErrorKind == input.ErrorKind && - contextRecord.ErrorCode == input.ErrorCode && - contextRecord.ErrorMessage == input.ErrorMessage && - sameJSON(contextRecord.ErrorDetails, input.ErrorDetails) && - contextRecord.RetryAt == nil && - contextRecord.Usage == usage && - contextRecord.ProviderReportedCostUSD == input.ProviderReportedCostUSD -} - -type modelCallErrorOutputInput struct { - ServedProviderModelSlug string - ErrorMessage string - Usage modelenvelope.Usage -} - -func publishModelCallErrorOutputTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - q *dbsqlc.Queries, - contextRow ModelCallContextRecord, - input modelCallErrorOutputInput, -) (TypedAgentEventRecord, error) { - authorityInput := modelCallErrorAuthorityInput(contextRow, input) - modelOutput, err := createModelOutputAuthorityTx(ctx, tx, authorityInput) - if err != nil { - return TypedAgentEventRecord{}, err - } - eventRecord, err := appendTypedAgentEventTx(ctx, txNotifications, tx, AppendTypedAgentEventInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - TurnID: modelOutput.TurnID, - Kind: events.KindModelOutput, - ModelOutputID: modelOutput.ID, - }) - if err != nil { - return TypedAgentEventRecord{}, err - } - if _, err := createContentBlockTx(ctx, tx, CreateContentBlockInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - OwnerKind: ContentBlockOwnerModelOutput, - OwnerModelOutputID: modelOutput.ID, - BlockKind: ContentBlockKindError, - TextContent: input.ErrorMessage, - }); err != nil { - return TypedAgentEventRecord{}, err - } - if err := updateAgentTurnLatestEventTx( - ctx, - tx, - contextRow.ProjectID, - contextRow.AgentID, - modelOutput.TurnID, - eventRecord.Event.ID, - eventRecord.Event.ID, - ); err != nil { - return TypedAgentEventRecord{}, err - } - if err := q.ReconcileAgentWakeup(ctx, dbsqlc.ReconcileAgentWakeupParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - Metadata: json.RawMessage(`{"reason":"model_call_error"}`), - }); err != nil { - return TypedAgentEventRecord{}, fmt.Errorf("reconcile wakeup after model call error: %w", err) - } - message := subagentMessage{ - Kind: SubagentMessageKindFailed, - Text: input.ErrorMessage, - IdempotencyKey: "model_output:" + modelOutput.ID.String(), - } - if err := handleSubagentTurnEndedTx( - ctx, txNotifications, tx, q, contextRow.ProjectID, contextRow.AgentID, message, - ); err != nil { - return TypedAgentEventRecord{}, err - } - return eventRecord, nil -} - -func modelCallErrorAuthorityInput( - contextRow ModelCallContextRecord, - input modelCallErrorOutputInput, -) CreateModelOutputAuthorityInput { - return CreateModelOutputAuthorityInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - ModelCallContextID: contextRow.ID, - ServedProviderModelSlug: input.ServedProviderModelSlug, - StopReason: modelenvelope.StopReasonError, - Usage: input.Usage, - } -} - -func replayModelCallErrorOutputTx( - ctx context.Context, - tx pgx.Tx, - q *dbsqlc.Queries, - contextRow ModelCallContextRecord, - input RecordModelCallErrorAndCompleteContextInput, -) (TypedAgentEventRecord, error) { - authorityInput := modelCallErrorAuthorityInput(contextRow, modelCallErrorOutputInput{ - ServedProviderModelSlug: input.ServedProviderModelSlug, - ErrorMessage: input.ErrorMessage, - Usage: input.Usage, - }) - modelOutput, err := createModelOutputAuthorityTx(ctx, tx, authorityInput) - if err != nil { - return TypedAgentEventRecord{}, err + return events.Event{}, err } - eventRecord, found, err := loadTypedEventByModelOutputMaybeTx( - ctx, - tx, - contextRow.ProjectID, - contextRow.AgentID, - modelOutput.ID, - ) + row, err := q.GetTypedAgentEventByModelOutput(ctx, dbsqlc.GetTypedAgentEventByModelOutputParams{ + ProjectID: input.ProjectID, AgentID: input.AgentID, ModelOutputID: &output.ID}) if err != nil { - return TypedAgentEventRecord{}, err - } - if !found || eventRecord.Event.Kind != events.KindModelOutput || - eventRecord.ModelOutputID != modelOutput.ID { - return TypedAgentEventRecord{}, storeerr.ErrIdempotencyConflict + return events.Event{}, err } - blocks, err := q.ListContentBlocksForModelOutput(ctx, dbsqlc.ListContentBlocksForModelOutputParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - ModelOutputID: storeutil.IDFromNil(modelOutput.ID), - }) + event, err := typedAgentEventFromModelOutputSQLC(row) if err != nil { - return TypedAgentEventRecord{}, fmt.Errorf("list terminal model error content blocks for replay: %w", err) + return events.Event{}, err } - if len(blocks) != 1 || blocks[0].Ordinal != 0 || blocks[0].BlockKind != string(ContentBlockKindError) || - blocks[0].TextContent != input.ErrorMessage || blocks[0].ArtifactID != nil || blocks[0].ToolCallID != nil { - return TypedAgentEventRecord{}, storeerr.ErrIdempotencyConflict + if err = unit.Commit(ctx, "record terminal model failure"); err != nil { + return events.Event{}, err } - return eventRecord, nil + return event.Event, nil } diff --git a/internal/storage/executionstore/model_outputs_workflow.go b/internal/storage/executionstore/model_outputs_workflow.go index 770d739ed..d486dc771 100644 --- a/internal/storage/executionstore/model_outputs_workflow.go +++ b/internal/storage/executionstore/model_outputs_workflow.go @@ -2,18 +2,11 @@ package executionstore import ( "context" - "errors" - "fmt" - "strings" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/events" "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" - "github.com/omnara-ai/omnara/internal/storage/storeerr" - "github.com/omnara-ai/omnara/internal/toolcatalog" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) type RecordModelOutputAndCompleteContextInput struct { @@ -22,10 +15,7 @@ type RecordModelOutputAndCompleteContextInput struct { RuntimeLockID uuid.UUID ModelCallContextID uuid.UUID ProviderRequestID string - // ProviderResponse is consumed inside the completion transaction and - // never durably stored. Storage and the model package share this type so - // the envelope shape is checked at compile time (no JSON re-parsing). - ProviderResponse modelenvelope.ResponseEnvelope + ProviderResponse modelenvelope.ResponseEnvelope } type ToolCallBindingInput struct { @@ -40,671 +30,76 @@ type RecordToolCallSourceAndCompleteContextInput struct { RuntimeLockID uuid.UUID ModelCallContextID uuid.UUID ProviderRequestID string - // ProviderResponse is consumed inside the completion transaction and - // never durably stored. Storage and the model package share this type so - // the envelope shape is checked at compile time (no JSON re-parsing). - ProviderResponse modelenvelope.ResponseEnvelope - ToolCallBindings []ToolCallBindingInput + ProviderResponse modelenvelope.ResponseEnvelope + ToolCallBindings []ToolCallBindingInput } -type boundToolCall struct { - ID uuid.UUID - ProviderCallID string - Name string - Input []byte - Type string - RejectionContent []byte -} - -func bindToolCalls( - envelope modelenvelope.ResponseEnvelope, +func executionOutput( + input RecordModelOutputAndCompleteContextInput, bindings []ToolCallBindingInput, -) ([]boundToolCall, error) { - if len(bindings) == 0 { - return nil, errors.New("tool call bindings are required") - } - byProviderCallID := make(map[string]ToolCallBindingInput, len(bindings)) - for _, binding := range bindings { - if binding.ProviderCallID == "" { - return nil, errors.New("tool call binding provider call id is required") - } - switch binding.Type { - case toolcatalog.ToolTypeBuiltIn, toolcatalog.ToolTypeCustom, toolcatalog.ToolTypeMCP: - default: - return nil, fmt.Errorf("unsupported tool call type %q", binding.Type) - } - if _, exists := byProviderCallID[binding.ProviderCallID]; exists { - return nil, fmt.Errorf( - "duplicate tool call binding for provider call id %q", - binding.ProviderCallID, - ) - } - byProviderCallID[binding.ProviderCallID] = binding - } - - toolCalls := make([]boundToolCall, 0, len(bindings)) - seenProviderCallIDs := make(map[string]struct{}, len(bindings)) - for _, part := range envelope.Normalized.Content { - if part.Type != modelenvelope.ResponsePartTypeToolCall { - continue - } - if _, exists := seenProviderCallIDs[part.ProviderCallID]; exists { - return nil, fmt.Errorf( - "duplicate provider call id %q in provider response", - part.ProviderCallID, - ) - } - if err := modelenvelope.ValidateToolInput(part.ToolInput); err != nil { - return nil, fmt.Errorf( - "provider call id %q input must be a JSON object", - part.ProviderCallID, - ) - } - seenProviderCallIDs[part.ProviderCallID] = struct{}{} - binding, ok := byProviderCallID[part.ProviderCallID] - if !ok { - return nil, fmt.Errorf( - "provider call id %q has no tool call binding", - part.ProviderCallID, - ) - } - call := boundToolCall{ - ID: binding.ID, - ProviderCallID: part.ProviderCallID, - Name: part.ToolName, - Input: part.ToolInput, - Type: binding.Type, - } - if part.ToolCallError != "" { - content, err := marshalJSON([]map[string]any{{ - "type": "structured_data", - "value": map[string]any{"error": part.ToolCallError, "error_code": "malformed"}, - }}) - if err != nil { - return nil, fmt.Errorf("marshal rejected tool call result: %w", err) - } - call.RejectionContent = content - } - toolCalls = append(toolCalls, call) - } - if len(toolCalls) == 0 { - return nil, errors.New("provider response contains no tool calls") - } - for _, binding := range bindings { - if _, used := seenProviderCallIDs[binding.ProviderCallID]; !used { - return nil, fmt.Errorf( - "tool call binding for provider call id %q has no provider response tool call", - binding.ProviderCallID, - ) - } - } - return toolCalls, nil -} - -func (s *Store) RecordToolCallSourceAndCompleteContext( - ctx context.Context, - input RecordToolCallSourceAndCompleteContextInput, -) (events.Event, []ToolCallRecord, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.ModelCallContextID == uuid.Nil { - return events.Event{}, nil, errors.New( - "project, agent, runtime lock, and model context are required", - ) - } - if err := validateModelResponseEnvelope(input.ProviderResponse); err != nil { - return events.Event{}, nil, err - } - toolCalls, err := bindToolCalls(input.ProviderResponse, input.ToolCallBindings) - if err != nil { - return events.Event{}, nil, err - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return events.Event{}, nil, fmt.Errorf("begin record tool call source: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { - return events.Event{}, nil, err - } - contextRow, err := loadModelCallContextByIDTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) - if err != nil { - return events.Event{}, nil, err - } - completionReplay, err := validateNormalModelCallCompletionState( - contextRow, - input.ModelCallContextID, - input.RuntimeLockID, - ) - if err != nil { - return events.Event{}, nil, err - } - if err := validateResponseEnvelopeForModelCallContext( - ctx, - dbsqlc.New(tx), - input.ProviderResponse, - contextRow, - ); err != nil { - return events.Event{}, nil, err - } - if completionReplay && !sameSuccessfulModelCallCompletionEvidence( - contextRow, - input.ProviderRequestID, - input.ProviderResponse, - ) { - return events.Event{}, nil, storeerr.ErrIdempotencyConflict - } - authorityInput := modelOutputAuthorityInputFromToolSourceContext(input) - modelOutput, err := createModelOutputAuthorityTx(ctx, tx, authorityInput) - if err != nil { - return events.Event{}, nil, err - } - eventRecord, foundExistingEvent, err := loadTypedEventByModelOutputMaybeTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.ID, - ) - if err != nil { - return events.Event{}, nil, err - } - if !foundExistingEvent { - eventRecord, err = appendTypedAgentEventTx( - ctx, - txNotifications, - tx, - AppendTypedAgentEventInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: modelOutput.TurnID, - Kind: events.KindModelOutput, - ModelOutputID: modelOutput.ID, - }, - ) - if err != nil { - return events.Event{}, nil, err - } - if err := updateAgentTurnLatestEventTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.TurnID, - eventRecord.Event.ID, - uuid.Nil, - ); err != nil { - return events.Event{}, nil, err - } - } - qtx := dbsqlc.New(tx) - if completionReplay { - if !foundExistingEvent { - return events.Event{}, nil, storeerr.ErrIdempotencyConflict - } - if err := validateModelOutputContentReplayTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.ID, - input.ProviderResponse, - ); err != nil { - return events.Event{}, nil, err - } - rows, err := qtx.ListToolCallsForModelContext( - ctx, - dbsqlc.ListToolCallsForModelContextParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ModelCallContextID: input.ModelCallContextID, - }, - ) - if err != nil { - return events.Event{}, nil, fmt.Errorf("list replayed tool calls: %w", err) - } - records := make([]ToolCallRecord, 0, len(rows)) - for _, row := range rows { - records = append(records, toolCallRecordFromContextSQLC(row)) - } - if !sameBoundToolCallBatch(records, toolCalls) { - return events.Event{}, nil, storeerr.ErrIdempotencyConflict - } - return eventRecord.Event, records, nil - } - if foundExistingEvent { - return events.Event{}, nil, storeerr.ErrIdempotencyConflict - } - - records := make([]ToolCallRecord, 0, len(toolCalls)) - for _, call := range toolCalls { - row, err := qtx.InsertToolCall(ctx, dbsqlc.InsertToolCallParams{ - ToolCallID: storeutil.IDFromNil(call.ID), - ProjectID: input.ProjectID, - AgentID: input.AgentID, - SourceEventID: eventRecord.Event.ID, - ModelCallContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ProviderCallID: call.ProviderCallID, - Name: call.Name, - Input: call.Input, - Type: call.Type, - }) - if err != nil { - if storeutil.IsUniqueViolation(err) { - return events.Event{}, nil, storeerr.ErrIdempotencyConflict - } - if errors.Is(err, pgx.ErrNoRows) { - if runtimeErr := agentRuntimeLockActiveTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); runtimeErr != nil { - return events.Event{}, nil, runtimeErr - } - return events.Event{}, nil, storeerr.ErrIdempotencyConflict - } - return events.Event{}, nil, fmt.Errorf("insert tool call proposal: %w", err) - } - record := toolCallRecordFromInsertSQLC(row) - records = append(records, record) - txNotifications.AddToolCallUpdate(record.AgentID, record.ID, string(record.State), nil) - } - toolCallContentBlockArgsByProviderCallID := make( - map[string]toolCallContentBlockArgs, - len(records), - ) - for _, record := range records { - toolCallContentBlockArgsByProviderCallID[record.ProviderCallID] = toolCallContentBlockArgs{ - ToolCallID: record.ID, - } - } - contentBlockByProviderCallID, err := createModelOutputContentBlocksTx( - ctx, - tx, - contextRow, - input.ProviderResponse, - modelOutput.ID, - toolCallContentBlockArgsByProviderCallID, - ) - if err != nil { - return events.Event{}, nil, err - } - for _, record := range records { - if contentBlockByProviderCallID[record.ProviderCallID] == uuid.Nil { - return events.Event{}, nil, fmt.Errorf( - "tool call %q has no matching tool_call content block in envelope content", - record.ProviderCallID, - ) - } +) agentexecution.AcceptOutputInput { + tools := make([]agentexecution.ToolProposal, len(bindings)) + for i, b := range bindings { + tools[i] = agentexecution.ToolProposal{ID: b.ID, ProviderCallID: b.ProviderCallID, Type: b.Type} + } + return agentexecution.AcceptOutputInput{ + RuntimeLockID: input.RuntimeLockID, + ContextID: input.ModelCallContextID, + RequestID: input.ProviderRequestID, + Response: input.ProviderResponse, + Tools: tools, } - for index, call := range toolCalls { - if len(call.RejectionContent) == 0 { - continue - } - record, err := completeToolCallTx(ctx, txNotifications, tx, CompleteToolCallInput{ - ProjectID: input.ProjectID, AgentID: input.AgentID, RuntimeLockID: input.RuntimeLockID, - ID: records[index].ID, Outcome: ToolResultOutcomeFailed, - ResultContentParts: call.RejectionContent, - }) - if err != nil { - return events.Event{}, nil, fmt.Errorf("record rejected tool call: %w", err) - } - records[index] = record - } - if err := completeSuccessfulNormalModelCallTx( - ctx, - qtx, - contextRow, - input.RuntimeLockID, - input.ProviderRequestID, - input.ProviderResponse, - ); err != nil { - return events.Event{}, nil, err - } - if err := qtx.MarkAgentWakeup(ctx, dbsqlc.MarkAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: []byte(`{"reason":"tool_work"}`), - }); err != nil { - return events.Event{}, nil, fmt.Errorf("mark tool work wakeup: %w", err) - } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "record tool call source"); err != nil { - return events.Event{}, nil, err - } - return eventRecord.Event, records, nil } -func sameBoundToolCallBatch( - existing []ToolCallRecord, - bound []boundToolCall, -) bool { - if len(existing) != len(bound) { - return false - } - existingByProviderCallID := make(map[string]ToolCallRecord, len(existing)) - for _, record := range existing { - existingByProviderCallID[record.ProviderCallID] = record - } - for _, call := range bound { - record, ok := existingByProviderCallID[call.ProviderCallID] - if !ok { - return false - } - if record.ProviderCallID != call.ProviderCallID || - record.Name != call.Name || - !sameJSON(record.Input, call.Input) || - record.Type != call.Type || - (call.ID != uuid.Nil && record.ID != call.ID) { - return false - } - if len(call.RejectionContent) > 0 && - (record.State != ToolCallStateCompleted || record.Outcome != ToolResultOutcomeFailed || - !sameJSON(record.ResultContentParts, call.RejectionContent)) { - return false - } - } - return true +func executionEvent(agentID uuid.UUID, event agentexecution.ExecutionEvent, kind events.Kind) events.Event { + return events.Event{ID: event.ID, AgentID: agentID, Sequence: event.Sequence, Kind: kind, At: event.Time} } func (s *Store) RecordModelOutputAndCompleteContext( ctx context.Context, input RecordModelOutputAndCompleteContextInput, ) (events.Event, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil || - input.ModelCallContextID == uuid.Nil { - return events.Event{}, errors.New( - "project, agent, runtime lock, and model context are required", - ) - } - if err := validateModelResponseEnvelope(input.ProviderResponse); err != nil { - return events.Event{}, err - } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return events.Event{}, fmt.Errorf("begin record model output: %w", err) - } - defer func() { _ = tx.Rollback(ctx) }() - if err := lockAgentWithParentTx(ctx, tx, dbsqlc.New(tx), input.ProjectID, input.AgentID); err != nil { return events.Event{}, err } - if err := ensureRuntimeLockActiveTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { - return events.Event{}, err - } - contextRow, err := loadModelCallContextByIDTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) - if err != nil { - return events.Event{}, err - } - completionReplay, err := validateNormalModelCallCompletionState( - contextRow, - input.ModelCallContextID, - input.RuntimeLockID, - ) + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.AcceptOutput(ctx, executionOutput(input, nil)) if err != nil { return events.Event{}, err } - if err := validateResponseEnvelopeForModelCallContext( - ctx, - dbsqlc.New(tx), - input.ProviderResponse, - contextRow, - ); err != nil { - return events.Event{}, err - } - if completionReplay && !sameSuccessfulModelCallCompletionEvidence( - contextRow, - input.ProviderRequestID, - input.ProviderResponse, - ) { - return events.Event{}, storeerr.ErrIdempotencyConflict - } - authorityInput := modelOutputAuthorityInputFromContext(input) - modelOutput, err := createModelOutputAuthorityTx(ctx, tx, authorityInput) - if err != nil { - return events.Event{}, err - } - if completionReplay { - if err := validateModelOutputContentReplayTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.ID, - input.ProviderResponse, - ); err != nil { - return events.Event{}, err - } - } - eventRecord, foundExistingEvent, err := loadTypedEventByModelOutputMaybeTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.ID, - ) - if err != nil { - return events.Event{}, err - } - if !foundExistingEvent { - eventRecord, err = appendTypedAgentEventTx( - ctx, - txNotifications, - tx, - AppendTypedAgentEventInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - TurnID: modelOutput.TurnID, - Kind: events.KindModelOutput, - ModelOutputID: modelOutput.ID, - }, - ) - if err != nil { - return events.Event{}, err - } - if err := updateAgentTurnLatestEventTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.TurnID, - eventRecord.Event.ID, - eventRecord.Event.ID, - ); err != nil { - return events.Event{}, err - } - } - if hasBlocks, err := modelOutputHasContentBlocksTx( - ctx, - tx, - input.ProjectID, - input.AgentID, - modelOutput.ID, - ); err != nil { - return events.Event{}, err - } else if !hasBlocks { - if _, err := createModelOutputContentBlocksTx( - ctx, - tx, - contextRow, - input.ProviderResponse, - modelOutput.ID, - nil, - ); err != nil { - return events.Event{}, err - } - } - if !completionReplay { - if err := completeSuccessfulNormalModelCallTx( - ctx, - dbsqlc.New(tx), - contextRow, - input.RuntimeLockID, - input.ProviderRequestID, - input.ProviderResponse, - ); err != nil { - return events.Event{}, err - } - if message, ended := subagentTurnEndMessage(input.ProviderResponse); ended { - message.IdempotencyKey = "model_output:" + modelOutput.ID.String() - if err := handleSubagentTurnEndedTx( - ctx, txNotifications, tx, dbsqlc.New(tx), input.ProjectID, input.AgentID, message, - ); err != nil { - return events.Event{}, err - } - } - } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "record model output"); err != nil { + if err = unit.Commit(ctx, "accept model output"); err != nil { return events.Event{}, err } - return eventRecord.Event, nil -} - -func validateModelResponseEnvelope(envelope modelenvelope.ResponseEnvelope) error { - return envelope.Validate() + return executionEvent(input.AgentID, result.Event, events.KindModelOutput), nil } -func validateResponseEnvelopeForModelCallContext( +func (s *Store) RecordToolCallSourceAndCompleteContext( ctx context.Context, - q *dbsqlc.Queries, - envelope modelenvelope.ResponseEnvelope, - contextRow ModelCallContextRecord, -) error { - requestedProviderModelSlug, err := q.GetModelCallContextProviderModelSlug( - ctx, - dbsqlc.GetModelCallContextProviderModelSlugParams{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - ModelCallContextID: contextRow.ID, - }, - ) + input RecordToolCallSourceAndCompleteContextInput, +) (events.Event, []ToolCallRecord, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { - return fmt.Errorf("load configured provider model slug: %w", err) - } - if envelope.RequestedProviderModelSlug != requestedProviderModelSlug { - return fmt.Errorf( - "provider response requested provider model slug %q does not match model call context %q", - envelope.RequestedProviderModelSlug, - requestedProviderModelSlug, - ) - } - if contextRow.State == ModelCallContextSucceeded && - (envelope.APIFormat != contextRow.APIFormat || envelope.APIVariant != contextRow.APIVariant) { - return fmt.Errorf( - "provider response route %q/%q does not match completed model call context %q/%q", - envelope.APIFormat, - envelope.APIVariant, - contextRow.APIFormat, - contextRow.APIVariant, - ) - } - return nil -} - -func validateNormalModelCallCompletionState( - contextRow ModelCallContextRecord, - modelCallContextID, runtimeLockID uuid.UUID, -) (bool, error) { - if contextRow.ID != modelCallContextID || - contextRow.OperationKind != ModelCallOperationNormal { - return false, storeerr.ErrStateTransitionConflict - } - if contextRow.RuntimeLockID != runtimeLockID { - return false, storeerr.ErrRuntimeLockInactive - } - if contextRow.State == ModelCallContextStarted { - return false, nil - } - if contextRow.State == ModelCallContextSucceeded { - return true, nil - } - return false, storeerr.ErrStateTransitionConflict -} - -func sameSuccessfulModelCallCompletionEvidence( - contextRow ModelCallContextRecord, - providerRequestID string, - envelope modelenvelope.ResponseEnvelope, -) bool { - return contextRow.ProviderRequestID == providerRequestID && - contextRow.ProviderReportedCostUSD == envelope.ProviderReportedCostUSD -} - -func completeSuccessfulNormalModelCallTx( - ctx context.Context, - q *dbsqlc.Queries, - contextRow ModelCallContextRecord, - runtimeLockID uuid.UUID, - providerRequestID string, - envelope modelenvelope.ResponseEnvelope, -) error { - if contextRow.OperationKind != ModelCallOperationNormal || - contextRow.State != ModelCallContextStarted { - return storeerr.ErrStateTransitionConflict + return events.Event{}, nil, err } - if _, err := finishModelCallContextTx(ctx, q, finishModelCallContextInput{ - ProjectID: contextRow.ProjectID, - AgentID: contextRow.AgentID, - ModelCallContextID: contextRow.ID, - RuntimeLockID: runtimeLockID, - ToState: ModelCallContextSucceeded, - APIFormat: envelope.APIFormat, - APIVariant: envelope.APIVariant, - ProviderRequestID: providerRequestID, - ProviderResponseID: envelope.Normalized.ID, - Usage: envelope.Normalized.Usage, - ProviderReportedCostUSD: envelope.ProviderReportedCostUSD, - ProviderMetadata: envelope.ProviderMetadata, - }); err != nil { - return err + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.AcceptOutput(ctx, executionOutput(RecordModelOutputAndCompleteContextInput{ + ProjectID: input.ProjectID, AgentID: input.AgentID, RuntimeLockID: input.RuntimeLockID, + ModelCallContextID: input.ModelCallContextID, ProviderRequestID: input.ProviderRequestID, + ProviderResponse: input.ProviderResponse}, input.ToolCallBindings)) + if err != nil { + return events.Event{}, nil, err } - return nil -} - -// subagentTurnEndMessage decides whether a model output ends a subagent's -// turn from the parent's point of view. Tool calls keep the turn open, and so -// does max_tokens, because the scheduler continues that output automatically; -// only end_turn delivers the answer, while refusal and content_filter tell the -// parent the subagent stopped without one. -func subagentTurnEndMessage(envelope modelenvelope.ResponseEnvelope) (subagentMessage, bool) { - if envelope.HasToolCalls() { - return subagentMessage{}, false + records := make([]ToolCallRecord, 0, len(result.Tools)) + for _, call := range result.Tools { + record, loadErr := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, call.ID) + if loadErr != nil { + return events.Event{}, nil, loadErr + } + records = append(records, record) } - text := strings.TrimSpace(envelope.Text()) - switch envelope.Normalized.StopReason { - case modelenvelope.StopReasonEndTurn: - return subagentMessage{Kind: SubagentMessageKindResult, Text: text}, true - case modelenvelope.StopReasonRefusal: - return subagentMessage{Kind: SubagentMessageKindRefused, Text: text}, true - case modelenvelope.StopReasonContentFilter: - return subagentMessage{Kind: SubagentMessageKindContentFiltered, Text: text}, true - default: - return subagentMessage{}, false + if err = unit.Commit(ctx, "accept model tool output"); err != nil { + return events.Event{}, nil, err } + return executionEvent(input.AgentID, result.Event, events.KindModelOutput), records, nil } diff --git a/internal/storage/executionstore/model_outputs_workflow_test.go b/internal/storage/executionstore/model_outputs_workflow_test.go deleted file mode 100644 index bafd66135..000000000 --- a/internal/storage/executionstore/model_outputs_workflow_test.go +++ /dev/null @@ -1,229 +0,0 @@ -package executionstore - -import ( - "context" - "encoding/json" - "strings" - "testing" - - "github.com/google/uuid" - "github.com/omnara-ai/omnara/internal/modelenvelope" - "github.com/omnara-ai/omnara/internal/toolcatalog" -) - -func TestValidateModelResponseEnvelopeRequiresStopReason(t *testing.T) { - envelope := modelenvelope.ResponseEnvelope{ - RequestedProviderModelSlug: "test-model", - APIFormat: "test-api", - APIVariant: "default", - Normalized: modelenvelope.ResponseNormalized{ - ID: "resp_missing_stop", - Content: []modelenvelope.ResponsePart{{Type: "text", Text: "hello"}}, - }, - } - err := validateModelResponseEnvelope(envelope) - if err == nil || !strings.Contains(err.Error(), "stop reason") { - t.Fatalf("validateModelResponseEnvelope error = %v, want missing stop reason error", err) - } -} - -func TestCreateModelOutputAuthorityRejectsUnsupportedStopReason(t *testing.T) { - for _, reason := range []modelenvelope.StopReason{ - modelenvelope.StopReasonPause, - modelenvelope.StopReasonContextWindow, - modelenvelope.StopReasonUnknown, - "not_a_stop_reason", - } { - t.Run(string(reason), func(t *testing.T) { - _, err := createModelOutputAuthorityTx(context.Background(), nil, CreateModelOutputAuthorityInput{ - ProjectID: uuid.MustParse("00000000-0000-4000-8000-000000000001"), - AgentID: uuid.MustParse("00000000-0000-4000-8000-000000000002"), - ModelCallContextID: uuid.MustParse("00000000-0000-4000-8000-000000000003"), - ServedProviderModelSlug: "test-model", - StopReason: reason, - }) - if err == nil || !strings.Contains(err.Error(), "unsupported model output stop reason") { - t.Fatalf("createModelOutputAuthorityTx error = %v, want unsupported stop reason error", err) - } - }) - } -} - -func TestBindToolCallsUsesBindingsForIdentityAndEnvelopeForContent(t *testing.T) { - firstID := uuid.MustParse("00000000-0000-4000-8000-000000000001") - secondID := uuid.MustParse("00000000-0000-4000-8000-000000000002") - envelope := validToolCallEnvelope( - modelenvelope.ResponsePart{Type: modelenvelope.ResponsePartTypeText, Text: "before tools"}, - modelenvelope.ResponsePart{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_first", - ToolName: "read_file", - ToolInput: json.RawMessage(`{"path":"README.md","nested":{"b":2,"a":1}}`), - }, - modelenvelope.ResponsePart{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_second", - ToolName: "docs__lookup", - ToolInput: json.RawMessage(`{"query":"durable agents"}`), - }, - ) - - bound, err := bindValidatedToolCalls( - envelope, - []ToolCallBindingInput{ - { - ID: secondID, - ProviderCallID: "call_second", - Type: toolcatalog.ToolTypeMCP, - }, - { - ID: firstID, - ProviderCallID: "call_first", - Type: toolcatalog.ToolTypeBuiltIn, - }, - }, - ) - if err != nil { - t.Fatalf("bind tool calls: %v", err) - } - if len(bound) != 2 { - t.Fatalf("bound tool calls = %+v, want two", bound) - } - if bound[0].ID != firstID || - bound[0].ProviderCallID != "call_first" || - bound[0].Type != toolcatalog.ToolTypeBuiltIn || - bound[0].Name != "read_file" || - !sameJSON(bound[0].Input, json.RawMessage(`{"nested":{"a":1,"b":2},"path":"README.md"}`)) { - t.Fatalf("first bound tool call = %+v", bound[0]) - } - if bound[1].ID != secondID || - bound[1].ProviderCallID != "call_second" || - bound[1].Type != toolcatalog.ToolTypeMCP || - bound[1].Name != "docs__lookup" || - !sameJSON(bound[1].Input, json.RawMessage(`{"query":"durable agents"}`)) { - t.Fatalf("second bound tool call = %+v", bound[1]) - } -} - -func TestBindToolCallsRejectsInconsistentBatch(t *testing.T) { - validEnvelope := validToolCallEnvelope(modelenvelope.ResponsePart{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_one", - ToolName: "read_file", - ToolInput: json.RawMessage(`{"path":"README.md"}`), - }) - validBinding := ToolCallBindingInput{ - ID: uuid.MustParse("00000000-0000-4000-8000-000000000001"), - ProviderCallID: "call_one", - Type: toolcatalog.ToolTypeBuiltIn, - } - - tests := []struct { - name string - envelope modelenvelope.ResponseEnvelope - bindings []ToolCallBindingInput - want string - }{ - { - name: "missing binding", - envelope: validEnvelope, - bindings: []ToolCallBindingInput{{ - ProviderCallID: "call_other", - Type: toolcatalog.ToolTypeBuiltIn, - }}, - want: "has no tool call binding", - }, - { - name: "extra binding", - envelope: validEnvelope, - bindings: []ToolCallBindingInput{ - validBinding, - { - ProviderCallID: "call_extra", - Type: toolcatalog.ToolTypeBuiltIn, - }, - }, - want: "has no provider response tool call", - }, - { - name: "duplicate binding", - envelope: validEnvelope, - bindings: []ToolCallBindingInput{ - validBinding, - validBinding, - }, - want: "duplicate tool call binding", - }, - { - name: "unsupported resolved type", - envelope: validEnvelope, - bindings: []ToolCallBindingInput{{ - ProviderCallID: "call_one", - Type: "unknown", - }}, - want: "unsupported tool call type", - }, - { - name: "invalid envelope input", - envelope: validToolCallEnvelope(modelenvelope.ResponsePart{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_one", - ToolName: "read_file", - ToolInput: json.RawMessage(`{"path":`), - }), - bindings: []ToolCallBindingInput{validBinding}, - want: "tool input must be a JSON object", - }, - { - name: "duplicate envelope provider identity", - envelope: validToolCallEnvelope( - validEnvelope.Normalized.Content[0], - validEnvelope.Normalized.Content[0], - ), - bindings: []ToolCallBindingInput{validBinding}, - want: "duplicate provider call id", - }, - { - name: "missing envelope tool name", - envelope: validToolCallEnvelope(modelenvelope.ResponsePart{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_one", - ToolInput: json.RawMessage(`{}`), - }), - bindings: []ToolCallBindingInput{validBinding}, - want: "missing provider_call_id or tool_name", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - _, err := bindValidatedToolCalls(tt.envelope, tt.bindings) - if err == nil || !strings.Contains(err.Error(), tt.want) { - t.Fatalf("bind tool calls error = %v, want %q", err, tt.want) - } - }) - } -} - -func validToolCallEnvelope(parts ...modelenvelope.ResponsePart) modelenvelope.ResponseEnvelope { - return modelenvelope.ResponseEnvelope{ - RequestedProviderModelSlug: "test-model", - APIFormat: "test-api", - APIVariant: "default", - Normalized: modelenvelope.ResponseNormalized{ - ID: "resp_tool_calls", - Content: parts, - StopReason: modelenvelope.StopReasonToolUse, - }, - } -} - -func bindValidatedToolCalls( - envelope modelenvelope.ResponseEnvelope, - bindings []ToolCallBindingInput, -) ([]boundToolCall, error) { - if err := validateModelResponseEnvelope(envelope); err != nil { - return nil, err - } - return bindToolCalls(envelope, bindings) -} diff --git a/internal/storage/executionstore/model_usage.go b/internal/storage/executionstore/model_usage.go index 86b9e86ce..6801e5f0d 100644 --- a/internal/storage/executionstore/model_usage.go +++ b/internal/storage/executionstore/model_usage.go @@ -1,45 +1,9 @@ package executionstore import ( - "math" - "github.com/omnara-ai/omnara/internal/modelenvelope" ) -type normalizedUsageColumns struct { - InputTokens *int - UncachedInputTokens *int - OutputTokens *int - ReasoningTokens *int - CacheReadTokens *int - CacheWriteTokens *int -} - -func modelUsageForStorage(usage modelenvelope.Usage) modelenvelope.Usage { - usage = modelenvelope.NormalizeUsage(usage) - if usage.InputTokens > math.MaxInt32 || - usage.UncachedInputTokens > math.MaxInt32 || - usage.OutputTokens > math.MaxInt32 || - usage.ReasoningTokens > math.MaxInt32 || - usage.CacheReadTokens > math.MaxInt32 || - usage.CacheWriteTokens > math.MaxInt32 { - return modelenvelope.Usage{} - } - return usage -} - -func usageColumnsFromModelUsage(usage modelenvelope.Usage) normalizedUsageColumns { - usage = modelUsageForStorage(usage) - return normalizedUsageColumns{ - InputTokens: modelenvelope.OptionalCount(usage.InputTokens), - UncachedInputTokens: modelenvelope.OptionalCount(usage.UncachedInputTokens), - OutputTokens: modelenvelope.OptionalCount(usage.OutputTokens), - ReasoningTokens: modelenvelope.OptionalCount(usage.ReasoningTokens), - CacheReadTokens: modelenvelope.OptionalCount(usage.CacheReadTokens), - CacheWriteTokens: modelenvelope.OptionalCount(usage.CacheWriteTokens), - } -} - func modelUsageFromSQLC( inputTokens, uncachedInputTokens, diff --git a/internal/storage/executionstore/model_usage_test.go b/internal/storage/executionstore/model_usage_test.go deleted file mode 100644 index 37b3c59eb..000000000 --- a/internal/storage/executionstore/model_usage_test.go +++ /dev/null @@ -1,57 +0,0 @@ -package executionstore - -import ( - "math" - "testing" - - "github.com/omnara-ai/omnara/internal/modelenvelope" -) - -func TestModelUsageForStorageNormalizesProviderUsage(t *testing.T) { - usage := modelUsageForStorage(modelenvelope.Usage{ - InputTokens: 20, - OutputTokens: 7, - ReasoningTokens: 3, - CacheReadTokens: 8, - CacheWriteTokens: 2, - }) - if usage != (modelenvelope.Usage{ - InputTokens: 20, - UncachedInputTokens: 10, - OutputTokens: 7, - ReasoningTokens: 3, - CacheReadTokens: 8, - CacheWriteTokens: 2, - }) { - t.Fatalf("normalized usage = %+v", usage) - } -} - -func TestModelUsageForStorageDropsInvalidOrUnrepresentableUsage(t *testing.T) { - tests := []modelenvelope.Usage{ - {InputTokens: -1}, - {InputTokens: math.MaxInt32 + 1}, - {InputTokens: 10, OutputTokens: 2, ReasoningTokens: 3}, - {InputTokens: 10, CacheReadTokens: 11}, - } - for _, usage := range tests { - if got := modelUsageForStorage(usage); got != (modelenvelope.Usage{}) { - t.Fatalf("invalid usage %+v normalized to %+v", usage, got) - } - } -} - -func TestUsageColumnsOmitUnreportedZeroValues(t *testing.T) { - columns := usageColumnsFromModelUsage(modelenvelope.Usage{ - InputTokens: 10, - UncachedInputTokens: 10, - OutputTokens: 4, - }) - if columns.InputTokens == nil || *columns.InputTokens != 10 || - columns.UncachedInputTokens == nil || *columns.UncachedInputTokens != 10 || - columns.OutputTokens == nil || *columns.OutputTokens != 4 || - columns.CacheReadTokens != nil || columns.CacheWriteTokens != nil || - columns.ReasoningTokens != nil { - t.Fatalf("usage columns = %+v", columns) - } -} diff --git a/internal/storage/executionstore/model_work_transition.go b/internal/storage/executionstore/model_work_transition.go new file mode 100644 index 000000000..bed2d5d62 --- /dev/null +++ b/internal/storage/executionstore/model_work_transition.go @@ -0,0 +1,57 @@ +package executionstore + +import ( + "context" +) + +type AcceptModelOutputAndAdvanceInput struct { + Output RecordModelOutputAndCompleteContextInput + ToolCallBindings []ToolCallBindingInput + AllowModelWork bool + PrepareModel bool +} + +type OwnedAgentWorkTransition struct { + Work ClaimedAgentWork + Continued bool +} + +func (s *Store) AcceptModelOutputAndAdvance( + ctx context.Context, + input AcceptModelOutputAndAdvanceInput, +) (OwnedAgentWorkTransition, error) { + output := input.Output + unit, h, err := beginExecution(ctx, s, output.ProjectID, output.AgentID) + if err != nil { + return OwnedAgentWorkTransition{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + accepted, work, err := h.AcceptAndAdvance( + ctx, + executionOutput(output, input.ToolCallBindings), + input.AllowModelWork, + input.PrepareModel, + ) + if err != nil { + return OwnedAgentWorkTransition{}, err + } + result := OwnedAgentWorkTransition{} + if accepted.Created { + if err := applyAdmissionDestination(ctx, + unit, + input.Output.ProjectID, + input.Output.AgentID, + work.Admission); err != nil { + return result, err + } + result.Work, err = shapeOwnedWork(ctx, unit, output.ProjectID, output.AgentID, work) + if err != nil { + return result, err + } + result.Continued = !work.Released + } + if err = unit.Commit(ctx, "accept model output and advance"); err != nil { + return OwnedAgentWorkTransition{}, err + } + return result, nil +} diff --git a/internal/storage/executionstore/model_work_transition_integration_test.go b/internal/storage/executionstore/model_work_transition_integration_test.go new file mode 100644 index 000000000..672207ac9 --- /dev/null +++ b/internal/storage/executionstore/model_work_transition_integration_test.go @@ -0,0 +1,214 @@ +//go:build integration + +package executionstore_test + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/modelprotocol" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +func fusedModelFixture( + t *testing.T, ctx context.Context, +) (processDaemonFixture, executionstore.AcceptModelOutputAndAdvanceInput) { + t.Helper() + f := newProcessDaemonFixture(t, ctx, "fused_output") + ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "initial") + advance := ownedAdvanceInput(f, true) + advance.PrepareModel = true + work, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, advance) + require.NoError(t, err) + require.True(t, continued) + require.NotNil(t, work.Model.Prepared) + require.True(t, work.Model.Prepared.Claim.Claimed) + model := work.Model.Prepared.Claim.Context + slug := modelProviderSlugForContext(t, ctx, f.Store, testProjectID, f.AgentID, model.ID) + return f, executionstore.AcceptModelOutputAndAdvanceInput{ + AllowModelWork: true, PrepareModel: true, + Output: executionstore.RecordModelOutputAndCompleteContextInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, ModelCallContextID: model.ID, + ProviderResponse: modelenvelope.ResponseEnvelope{ + RequestedProviderModelSlug: slug, ServedProviderModelSlug: slug, + APIFormat: modelprotocol.APIFormatOpenAIResponses, APIVariant: "default", + Normalized: modelenvelope.ResponseNormalized{ID: "fused_response", StopReason: modelenvelope.StopReasonEndTurn}, + }, + }, + } +} + +func TestAcceptModelOutputAndAdvancePrecedenceAndPreparation(t *testing.T) { + for _, mode := range []string{"idle", "queued", "steering", "output limit", "budget", "tools before steering"} { + t.Run(mode, func(t *testing.T) { + ctx := t.Context() + f, input := fusedModelFixture(t, ctx) + var queued, steering executionstore.AgentInputRecord + if mode != "idle" { + queued = ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "queued") + } + if mode == "steering" || mode == "tools before steering" { + steering = ownedInput(t, ctx, f, executionstore.DeliveryModeSteering, "steering") + } + if mode == "output limit" || mode == "steering" { + input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonMaxTokens + } + if mode == "budget" { + input.AllowModelWork = false + } + if mode == "tools before steering" { + input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonToolUse + input.Output.ProviderResponse.Normalized.Content = []modelenvelope.ResponsePart{{ + Type: modelenvelope.ResponsePartTypeToolCall, ProviderCallID: "call", + ToolName: "read_process", ToolInput: []byte(`{}`), + }} + input.ToolCallBindings = []executionstore.ToolCallBindingInput{{ProviderCallID: "call", Type: "built_in"}} + } + transition, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.NoError(t, err) + wantContinued := mode != "idle" && mode != "budget" + require.Equal(t, wantContinued, transition.Continued) + requireOwnedLease(t, ctx, f, wantContinued) + completed, found, err := f.Store.Execution().GetModelCallContext( + ctx, testProjectID, f.AgentID, input.Output.ModelCallContextID, + ) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, executionstore.ModelCallContextSucceeded, completed.State) + if !wantContinued { + return + } + work := transition.Work + require.Equal(t, f.Lock.ID, work.RuntimeLock.ID) + require.Zero(t, countAgentWakeups(t, ctx, f.Store, f.AgentID)) + if mode == "tools before steering" { + require.Equal(t, executionstore.AgentWorkTool, work.Kind) + require.NotNil(t, work.Tool.Prepared) + require.NotNil(t, work.Tool.Prepared.FirstCall) + require.Equal(t, "call", work.Tool.Prepared.FirstCall.ProviderCallID) + require.Equal(t, executionstore.ToolCallStateAwaitingAuthorization, work.Tool.Prepared.FirstCall.State) + return + } + require.Equal(t, executionstore.AgentWorkModel, work.Kind) + prepared := work.Model.Prepared + require.NotNil(t, prepared) + require.True(t, prepared.Claim.Claimed) + require.NotNil(t, prepared.ContextData) + require.Equal(t, prepared.Claim.Context.InputEventSequence, prepared.ContextData.InputEventSequence) + require.NotEmpty(t, prepared.ContextData.Events) + switch mode { + case "queued": + require.Equal(t, []uuid.UUID{queued.ID}, work.Model.InputIDs) + case "steering": + require.Equal(t, []uuid.UUID{steering.ID}, work.Model.InputIDs) + case "output limit": + require.Equal(t, executionstore.ModelWorkContinue, work.Model.Kind) + require.Equal(t, input.Output.ModelCallContextID, work.Model.SourceModelCallContextID) + } + }) + } +} + +func TestAcceptModelOutputAndAdvanceFencing(t *testing.T) { + for _, mode := range []string{"cancel", "expired", "replaced"} { + t.Run(mode, func(t *testing.T) { + ctx := t.Context() + f, input := fusedModelFixture(t, ctx) + if mode == "cancel" { + _, err := f.Store.Execution().CancelAgent(ctx, executionstore.CancelAgentInput{ + ProjectID: testProjectID, AgentID: f.AgentID, Actor: mustOmnaraActorParams(t, f.UserID), + }) + require.NoError(t, err) + } else { + expireAgentRuntimeLockForTest(t, ctx, f.Store, f.Lock.ID) + if mode == "replaced" { + n, err := f.Store.Execution().ReapExpiredAgentRuntimeLocks(ctx, 100) + require.NoError(t, err) + require.EqualValues(t, 1, n) + _, found, err := f.Store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + } + } + _, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) + _, found, err := f.Store.Execution().GetModelOutputForContext( + ctx, testProjectID, f.AgentID, input.Output.ModelCallContextID, + ) + require.NoError(t, err) + require.False(t, found) + }) + } +} + +func TestAcceptModelOutputAndAdvanceRollbackAndCrashRecovery(t *testing.T) { + ctx := t.Context() + f, input := fusedModelFixture(t, ctx) + input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonMaxTokens + _, err := f.Store.pool.Exec(ctx, `CREATE FUNCTION reject_fused_preparation() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN RAISE EXCEPTION 'reject fused preparation'; END $$; + CREATE CONSTRAINT TRIGGER reject_fused_preparation AFTER INSERT ON model_call_contexts + DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION reject_fused_preparation();`) + require.NoError(t, err) + _, err = f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.ErrorContains(t, err, "commit accept model output and advance") + _, found, err := f.Store.Execution().GetModelOutputForContext( + ctx, testProjectID, f.AgentID, input.Output.ModelCallContextID, + ) + require.NoError(t, err) + require.False(t, found) + original, found, err := f.Store.Execution().GetModelCallContext( + ctx, testProjectID, f.AgentID, input.Output.ModelCallContextID, + ) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, executionstore.ModelCallContextStarted, original.State) + requireOwnedLease(t, ctx, f, true) + _, err = f.Store.pool.Exec(ctx, `DROP TRIGGER reject_fused_preparation ON model_call_contexts`) + require.NoError(t, err) + transition, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.NoError(t, err) + require.True(t, transition.Continued) + next := transition.Work.Model.Prepared.Claim.Context + expireAgentRuntimeLockForTest(t, ctx, f.Store, f.Lock.ID) + n, err := f.Store.Execution().ReapExpiredAgentRuntimeLocks(ctx, 100) + require.NoError(t, err) + require.EqualValues(t, 1, n) + recovered, found, err := f.Store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, executionstore.ModelWorkResume, recovered.Model.Kind) + require.Equal(t, next.ID, recovered.Model.ModelCallContextID) + require.NotEqual(t, f.Lock.ID, recovered.RuntimeLock.ID) + input.Output.ModelCallContextID = next.ID + _, err = f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) +} + +func TestAcceptModelOutputAndAdvanceReplayRecoversPreparedWorkWithoutDuplicatingOutput(t *testing.T) { + ctx := t.Context() + f, input := fusedModelFixture(t, ctx) + input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonMaxTokens + transition, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.NoError(t, err) + require.True(t, transition.Continued) + next := transition.Work.Model.Prepared.Claim.Context + replay, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + require.NoError(t, err) + require.False(t, replay.Continued) + requireOwnedLease(t, ctx, f, false) + recovered, found, err := f.Store.Execution().GetModelCallContext(ctx, testProjectID, f.AgentID, next.ID) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, executionstore.ModelCallContextFailed, recovered.State) + require.Equal(t, executionstore.ModelCallRecoveryRetry, recovered.RecoveryKind) + var outputs int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM model_outputs WHERE agent_id=$1`, f.AgentID, + ).Scan(&outputs)) + require.Equal(t, 1, outputs) +} diff --git a/internal/storage/executionstore/org_lifecycle_ops.go b/internal/storage/executionstore/org_lifecycle_ops.go index 0ecfa6ae6..cbc15b08b 100644 --- a/internal/storage/executionstore/org_lifecycle_ops.go +++ b/internal/storage/executionstore/org_lifecycle_ops.go @@ -8,29 +8,30 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) -func (s *Store) ArchiveAgentTx( +func (s *Store) ArchiveAgentInUnit( ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, + unit *agentexecution.Unit, projectID, agentID uuid.UUID, actor *ActorParams, ) ([]MachineRecord, error) { - return archiveAgentTreeTx(ctx, tx, s.q.WithTx(tx), txNotifications, projectID, agentID, actor) + + tx := unit.DB() + return archiveAgentTreeTx(ctx, unit, dbsqlc.New(tx), projectID, agentID, actor) } func (s *Store) ProvisionOrganizationDefaultsTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, orgID, projectID uuid.UUID, templates []DefaultMachinePoolTemplate, ) error { - qtx := s.q.WithTx(tx) + qtx := dbsqlc.New(tx) for _, template := range templates { input := template.createInput(orgID) poolDefaults, err := prepareMachinePoolCreateInput(&input) diff --git a/internal/storage/executionstore/pool_machines_integration_test.go b/internal/storage/executionstore/pool_machines_integration_test.go index fb0ad107b..862be8a58 100644 --- a/internal/storage/executionstore/pool_machines_integration_test.go +++ b/internal/storage/executionstore/pool_machines_integration_test.go @@ -15,11 +15,11 @@ import ( "github.com/omnara-ai/omnara/internal/machinepool" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/modelprotocol" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/secrets" "github.com/omnara-ai/omnara/internal/storage" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/secretstore" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/toolcatalog" @@ -47,12 +47,14 @@ func TestListMachinePoolSourcesUsesCurrentNamesAfterSwap(t *testing.T) { 2, now.Add(time.Duration(index)*time.Second), ) - if _, err := store.Execution().CreateProjectMachinePoolGrant(ctx, executionstore.CreateProjectMachinePoolGrantInput{ - OrgID: testOrgID, - ProjectID: testProjectID, - MachinePoolID: machinePools[index].ID, - IdempotencyKey: fmt.Sprintf("idem-agent-pool-name-swap-grant-%d", index), - }); err != nil { + if _, + err := store.Execution().CreateProjectMachinePoolGrant(ctx, + executionstore.CreateProjectMachinePoolGrantInput{ + OrgID: testOrgID, + ProjectID: testProjectID, + MachinePoolID: machinePools[index].ID, + IdempotencyKey: fmt.Sprintf("idem-agent-pool-name-swap-grant-%d", index), + }); err != nil { t.Fatalf("create pool grant: %v", err) } } @@ -285,7 +287,8 @@ func TestCreatePoolMachineUsesCurrentSourceWhilePoolRemainsConfigured(t *testing if err != nil { t.Fatalf("load rolled-back tool call: %v", err) } - if rolledBackToolCall.State != executionstore.ToolCallStateReady || rolledBackToolCall.RuntimeLockID != uuid.Nil { + if rolledBackToolCall.State != executionstore.ToolCallStateReady || + rolledBackToolCall.RuntimeLockID != uuid.Nil { t.Fatalf( "rolled-back tool call state/runtime = %s/%s, want ready/unowned", rolledBackToolCall.State, @@ -324,7 +327,10 @@ func TestCreatePoolMachineUsesCurrentSourceWhilePoolRemainsConfigured(t *testing result.Machine.Machine.ProviderOptions, json.RawMessage(`{"image":"captured","startup_script":"current"}`), ) { - t.Fatalf("created machine provider options = %s, want current provisioning", result.Machine.Machine.ProviderOptions) + t.Fatalf( + "created machine provider options = %s, want current provisioning", + result.Machine.Machine.ProviderOptions, + ) } staleTransaction := executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, @@ -362,8 +368,17 @@ func TestLaunchAndCreatePoolMachineAfterEnvironmentSecretDeletion(t *testing.T) t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "deleted_machine_secret") - store := newIntegrationStore(fixture.Store.pool, storage.WithMachinePoolProviders(mergingMachinePoolProviders{})) - user := mustCreateProjectDeveloperUser(t, ctx, store, "deleted-machine-secret@example.com", "Secret Tester") + store := newIntegrationStore( + fixture.Store.pool, + storage.WithMachinePoolProviders(mergingMachinePoolProviders{}), + ) + user := mustCreateProjectDeveloperUser( + t, + ctx, + store, + "deleted-machine-secret@example.com", + "Secret Tester", + ) secret, _, err := store.Secrets().CreateSecret(ctx, secretstore.CreateSecretInput{ OrgID: testOrgID, OwnerKind: secretstore.SecretOwnerProject, @@ -383,14 +398,16 @@ func TestLaunchAndCreatePoolMachineAfterEnvironmentSecretDeletion(t *testing.T) DefaultMachineEnv: json.RawMessage(`{"PLAIN":"plain"}`), DefaultMachineProviderOptions: json.RawMessage(`{"image":"test"}`), }, 2, fixture.Now) - if _, err := store.Execution().CreateProjectMachinePoolGrant(ctx, projectGrantInputWithDefaultMachineOverlayForTest( - executionstore.CreateProjectMachinePoolGrantInput{ - OrgID: testOrgID, ProjectID: testProjectID, MachinePoolID: machinePool.ID, - }, - defaultMachineOverlayFieldsForTest{ - DefaultMachineSecretEnvOverlay: json.RawMessage(`{"GRANT_SECRET":"` + secret.ID.String() + `"}`), - }, - )); err != nil { + if _, + err := store.Execution().CreateProjectMachinePoolGrant(ctx, + projectGrantInputWithDefaultMachineOverlayForTest( + executionstore.CreateProjectMachinePoolGrantInput{ + OrgID: testOrgID, ProjectID: testProjectID, MachinePoolID: machinePool.ID, + }, + defaultMachineOverlayFieldsForTest{ + DefaultMachineSecretEnvOverlay: json.RawMessage(`{"GRANT_SECRET":"` + secret.ID.String() + `"}`), + }, + )); err != nil { t.Fatalf("grant machine pool: %v", err) } machine, err := store.Execution().GetMachine(ctx, testOrgID, fixture.MachineID) @@ -439,8 +456,17 @@ tools: if err != nil { t.Fatalf("acquire agent runtime: %v", err) } - toolCalls := createPoolMachineToolCalls(t, ctx, store, launch.Agent.ID, user.ID, config.ID, lock, - "deleted_secret", []poolMachineToolCallSpec{{Label: "create", Name: "create_machine", Input: json.RawMessage(`{}`)}}) + toolCalls := createPoolMachineToolCalls( + t, + ctx, + store, + launch.Agent.ID, + user.ID, + config.ID, + lock, + "deleted_secret", + []poolMachineToolCallSpec{{Label: "create", Name: "create_machine", Input: json.RawMessage(`{}`)}}, + ) created, err := createPoolMachineForTest(ctx, store, executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, AgentID: launch.Agent.ID, ToolCallID: toolCalls["create"], RuntimeLockID: lock.ID, }, executionstore.CreatePoolMachineInput{MachinePoolID: machinePool.ID}) @@ -498,9 +524,11 @@ func TestCreatePoolMachineUsesResolvedConfigAndCwd(t *testing.T) { MaxTotalMachines: 3, }, defaultMachineFieldsForTest{ - DefaultMachineCPU: 4, - DefaultMachineMemoryMB: 8192, - DefaultMachineEnv: json.RawMessage(`{"POOL":"base","SHARED":"pool","REMOVE":"pool"}`), + DefaultMachineCPU: 4, + DefaultMachineMemoryMB: 8192, + DefaultMachineEnv: json.RawMessage( + `{"POOL":"base","SHARED":"pool","REMOVE":"pool"}`, + ), DefaultMachineProviderOptions: json.RawMessage(`{"image":"pool","pool_only":"pool"}`), }, )), @@ -509,20 +537,22 @@ func TestCreatePoolMachineUsesResolvedConfigAndCwd(t *testing.T) { if err != nil { t.Fatalf("create machine pool: %v", err) } - if _, err := store.Execution().CreateProjectMachinePoolGrant(ctx, projectGrantInputWithDefaultMachineOverlayForTest( - executionstore.CreateProjectMachinePoolGrantInput{ - OrgID: testOrgID, - ProjectID: testProjectID, - MachinePoolID: machinePool.ID, - DefaultCwd: "/grant", - IdempotencyKey: "idem-agent-pool-machine-resolved-grant", - }, - defaultMachineOverlayFieldsForTest{ - DefaultMachineMemoryMB: new(4096), - DefaultMachineEnvOverlay: json.RawMessage(`{"GRANT":"one","SHARED":"grant","REMOVE":null}`), - DefaultMachineProviderOptionsOverlay: json.RawMessage(`{"image":"grant","grant_only":"grant"}`), - }, - )); err != nil { + if _, + err := store.Execution().CreateProjectMachinePoolGrant(ctx, + projectGrantInputWithDefaultMachineOverlayForTest( + executionstore.CreateProjectMachinePoolGrantInput{ + OrgID: testOrgID, + ProjectID: testProjectID, + MachinePoolID: machinePool.ID, + DefaultCwd: "/grant", + IdempotencyKey: "idem-agent-pool-machine-resolved-grant", + }, + defaultMachineOverlayFieldsForTest{ + DefaultMachineMemoryMB: new(4096), + DefaultMachineEnvOverlay: json.RawMessage(`{"GRANT":"one","SHARED":"grant","REMOVE":null}`), + DefaultMachineProviderOptionsOverlay: json.RawMessage(`{"image":"grant","grant_only":"grant"}`), + }, + )); err != nil { t.Fatalf("create pool grant: %v", err) } config := mustCreateAgentConfigFromYAML( @@ -613,7 +643,9 @@ func TestCreatePoolMachineUsesResolvedConfigAndCwd(t *testing.T) { requireMachineEnvironmentForTest( t, machineEnvironmentFromRecordForTest(t, created.Machine.Machine), - executionstore.MachineEnvironment{Env: map[string]string{"POOL": "base", "GRANT": "one", "SHARED": "grant"}}, + executionstore.MachineEnvironment{ + Env: map[string]string{"POOL": "base", "GRANT": "one", "SHARED": "grant"}, + }, ) cwd := "/mutated" env := json.RawMessage(`{"MUTATED":"true"}`) @@ -1333,7 +1365,10 @@ func TestCreatePoolMachineReplayMaxAndDeleteLifecycle(t *testing.T) { ) } if created.Machine.Binding.DeleteToolCallID != uuid.Nil { - t.Fatalf("created machine binding delete tool call = %s, want nil", created.Machine.Binding.DeleteToolCallID) + t.Fatalf( + "created machine binding delete tool call = %s, want nil", + created.Machine.Binding.DeleteToolCallID, + ) } if created.Machine.Machine.IdempotencyKey != "" { t.Fatalf("created machine idempotency key = %q, want empty", created.Machine.Machine.IdempotencyKey) @@ -1400,7 +1435,8 @@ func TestCreatePoolMachineReplayMaxAndDeleteLifecycle(t *testing.T) { if err != nil { t.Fatalf("delete pool machine: %v", err) } - if deleted.Machine.LifecycleState != "deleting" || deleted.Machine.LifecycleReasonCode != "machine_tool_delete" { + if deleted.Machine.LifecycleState != "deleting" || + deleted.Machine.LifecycleReasonCode != "machine_tool_delete" { t.Fatalf("deleted machine = %+v", deleted.Machine) } if deleted.Binding.CreateToolCallID != createTransaction.ToolCallID { @@ -1487,13 +1523,14 @@ func TestCreatePoolMachineReplayMaxAndDeleteLifecycle(t *testing.T) { t.Fatalf("replacement pool machine while deleting error = %v, want state transition conflict", err) } - claimedDelete, ok, err := store.Execution().ClaimPoolMachineDeletion(ctx, executionstore.MachineDeletingInput{ - OrgID: testOrgID, - MachineID: created.Machine.Machine.ID, - LifecycleReasonCode: cleanup[0].ReasonCode, - LifecycleReasonMessage: cleanup[0].ReasonMessage, - ExpectedLifecycleVersion: cleanup[0].Machine.LifecycleVersion, - }) + claimedDelete, ok, err := store.Execution(). + ClaimPoolMachineDeletion(ctx, executionstore.MachineDeletingInput{ + OrgID: testOrgID, + MachineID: created.Machine.Machine.ID, + LifecycleReasonCode: cleanup[0].ReasonCode, + LifecycleReasonMessage: cleanup[0].ReasonMessage, + ExpectedLifecycleVersion: cleanup[0].Machine.LifecycleVersion, + }) if err != nil || !ok { t.Fatalf("claim pool machine deletion ok=%v err=%v", ok, err) } @@ -1517,7 +1554,10 @@ func TestCreatePoolMachineReplayMaxAndDeleteLifecycle(t *testing.T) { if _, err := createPoolMachineForTest( ctx, store, replacementTransaction, replacementInput, ); !errors.Is(err, storeerr.ErrStateTransitionConflict) { - t.Fatalf("replacement pool machine while delete_failed error = %v, want state transition conflict", err) + t.Fatalf( + "replacement pool machine while delete_failed error = %v, want state transition conflict", + err, + ) } } @@ -1542,7 +1582,9 @@ func TestDeletePoolMachineAllowsFreshProvisioningMachine(t *testing.T) { store, "Provisioning Delete Pool", "test.provider", - defaultMachineFieldsForTest{DefaultMachineProviderOptions: json.RawMessage(`{"image":"provisioning-delete"}`)}, + defaultMachineFieldsForTest{ + DefaultMachineProviderOptions: json.RawMessage(`{"image":"provisioning-delete"}`), + }, 3, now, ) @@ -1630,7 +1672,8 @@ func TestDeletePoolMachineAllowsFreshProvisioningMachine(t *testing.T) { if err != nil { t.Fatalf("delete provisioning pool machine: %v", err) } - if deleted.Machine.LifecycleState != "deleting" || deleted.Machine.LifecycleReasonCode != "machine_tool_delete" { + if deleted.Machine.LifecycleState != "deleting" || + deleted.Machine.LifecycleReasonCode != "machine_tool_delete" { t.Fatalf("deleted provisioning machine = %+v", deleted.Machine) } if err := store.Execution().CompletePoolMachineProvisioning( @@ -1960,7 +2003,8 @@ func TestPoolMachineToolsExcludeExplicitPoolBackedMachineSource(t *testing.T) { if err != nil { t.Fatalf("load deleted pool-backed machine: %v", err) } - if deletedMachine.LifecycleState != executionstore.MachineLifecycleStateDeleted || deletedMachine.DeletedAt == nil { + if deletedMachine.LifecycleState != executionstore.MachineLifecycleStateDeleted || + deletedMachine.DeletedAt == nil { t.Fatalf("deleted pool-backed machine = %+v", deletedMachine) } if count := countProjectMachineGrantsForMachineForTest( @@ -2129,16 +2173,15 @@ func activateAgentConfigForPoolMachineTest( idempotencyKey string, ) { t.Helper() - tx, err := store.pool.Begin(ctx) + tx, err := store.Execution().IntegrationBeginUnit(ctx) if err != nil { t.Fatalf("begin activate pool machine config: %v", err) } defer func() { _ = tx.Rollback(ctx) }() if _, err := executionstore.IntegrationActivateAgentConfigTx( ctx, - notifications.NewTxNotifications(), tx, - store.q.WithTx(tx), + dbsqlc.New(tx.DB()), executionstore.ActivateAgentConfigInput{ ProjectID: testProjectID, AgentID: agentID, @@ -2150,7 +2193,7 @@ func activateAgentConfigForPoolMachineTest( ); err != nil { t.Fatalf("activate pool machine config: %v", err) } - if err := tx.Commit(ctx); err != nil { + if err := tx.Commit(ctx, "config activation"); err != nil { t.Fatalf("commit activate pool machine config: %v", err) } } @@ -2173,12 +2216,13 @@ func createPoolMachineToolCalls( Actor: mustOmnaraActorParams(t, userID), ContentBlocks: json.RawMessage(`[{"type":"text","text":"seed machine tool calls"}]`), IdempotencyKey: "pool-machine-tool-input-" + label, + DeliveryMode: executionstore.DeliveryModeSteering, }, ) if err != nil { t.Fatalf("create tool-call seed input: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, store, @@ -2189,17 +2233,17 @@ func createPoolMachineToolCalls( if !found { t.Fatal("expected tool-call seed input admission") } - claim, err := store.Execution().ClaimNormalModelCall( + prepared1, err := store.Execution().PrepareNormalModelCall( ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: agentID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: configID, - InputEventSequence: admitted.Events[0].Sequence, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: agentID, + RuntimeLockID: lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, }, ) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim tool-call seed model context: %v", err) } @@ -2306,16 +2350,28 @@ func TestCreatePoolMachineSizing(t *testing.T) { seedMigratedDB(t, ctx, pool) store := newIntegrationStore(pool, storage.WithMachinePoolProviders(machinepool.DefaultCatalog())) user := mustCreateProjectDeveloperUser(t, ctx, store, "sizing@example.com", "Sizing") - machinePool := createLaunchTestMachinePool(t, ctx, store, "Sizing Pool", "unikraft", defaultMachineFieldsForTest{ - DefaultMachineCPU: 1, DefaultMachineMemoryMB: 1024, - DefaultMachineProviderOptions: json.RawMessage(`{"image":"test-image","metro":"sfo"}`), - }, 8, time.Now()) - _, err := store.Execution().CreateProjectMachinePoolGrant(ctx, executionstore.CreateProjectMachinePoolGrantInput{ - OrgID: testOrgID, ProjectID: testProjectID, MachinePoolID: machinePool.ID, IdempotencyKey: "sizing", - DefaultMachineCPU: new(2), DefaultMachineMemoryMB: new(4096), - MinMachineCPU: new(2), MaxMachineCPU: new(4), MinMachineMemoryMB: new(2048), MaxMachineMemoryMB: new(8192), - MaxTotalCPU: new(12), MaxTotalMemoryMB: new(24576), - }) + machinePool := createLaunchTestMachinePool( + t, + ctx, + store, + "Sizing Pool", + "unikraft", + defaultMachineFieldsForTest{ + DefaultMachineCPU: 1, DefaultMachineMemoryMB: 1024, + DefaultMachineProviderOptions: json.RawMessage(`{"image":"test-image","metro":"sfo"}`), + }, + 8, + time.Now(), + ) + _, err := store.Execution(). + CreateProjectMachinePoolGrant(ctx, executionstore.CreateProjectMachinePoolGrantInput{ + OrgID: testOrgID, ProjectID: testProjectID, MachinePoolID: machinePool.ID, IdempotencyKey: "sizing", + DefaultMachineCPU: new(2), DefaultMachineMemoryMB: new(4096), + MinMachineCPU: new( + 2, + ), MaxMachineCPU: new(4), MinMachineMemoryMB: new(2048), MaxMachineMemoryMB: new(8192), + MaxTotalCPU: new(12), MaxTotalMemoryMB: new(24576), + }) require.NoError(t, err) config := mustCreateAgentConfigFromYAML( t, ctx, store, @@ -2350,12 +2406,20 @@ func TestCreatePoolMachineSizing(t *testing.T) { {name: "memory only", memory: new(8192), wantCPU: 3, wantMemory: 8192}, {name: "both", cpu: new(2), memory: new(6144), wantCPU: 2, wantMemory: 6144}, {name: "pool maximum", cpu: new(33), wantError: "machine pool per-machine cpu capacity exceeded"}, - {name: "project cpu maximum", cpu: new(5), wantError: "project machine pool per-machine cpu capacity exceeded"}, + { + name: "project cpu maximum", + cpu: new(5), + wantError: "project machine pool per-machine cpu capacity exceeded", + }, { name: "project memory maximum", memory: new(16384), wantError: "project machine pool per-machine memory capacity exceeded", }, - {name: "project cpu minimum", cpu: new(1), wantError: "project machine pool per-machine cpu minimum not met"}, + { + name: "project cpu minimum", + cpu: new(1), + wantError: "project machine pool per-machine cpu minimum not met", + }, { name: "project memory minimum", memory: new(1024), wantError: "project machine pool per-machine memory minimum not met", @@ -2364,7 +2428,10 @@ func TestCreatePoolMachineSizing(t *testing.T) { } specs := make([]poolMachineToolCallSpec, 0, len(tests)) for _, test := range tests { - specs = append(specs, poolMachineToolCallSpec{Label: test.name, Name: "create_machine", Input: json.RawMessage(`{}`)}) + specs = append( + specs, + poolMachineToolCallSpec{Label: test.name, Name: "create_machine", Input: json.RawMessage(`{}`)}, + ) } calls := createPoolMachineToolCalls(t, ctx, store, agent.ID, user.ID, config.ID, lock, "sizing", specs) for _, test := range tests { @@ -2372,7 +2439,11 @@ func TestCreatePoolMachineSizing(t *testing.T) { transaction := executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, AgentID: agent.ID, ToolCallID: calls[test.name], RuntimeLockID: lock.ID, } - input := executionstore.CreatePoolMachineInput{MachinePoolID: machinePool.ID, CPU: test.cpu, MemoryMB: test.memory} + input := executionstore.CreatePoolMachineInput{ + MachinePoolID: machinePool.ID, + CPU: test.cpu, + MemoryMB: test.memory, + } created, err := createPoolMachineForTest(ctx, store, transaction, input) if test.wantError != "" { require.ErrorContains(t, err, test.wantError) @@ -2414,12 +2485,19 @@ func TestCreatePoolMachineRejectsUnsupportedSizing(t *testing.T) { user := mustCreateProjectDeveloperUser(t, ctx, store, "unsupported@example.com", "Sizing") input := executionstore.CreateMachinePoolInput{ OrgID: testOrgID, Name: "Sizing Pool", Provider: provider, MaxTotalMachines: 4, - ProviderAuthSecretID: createMachinePoolProviderAuthSecretForTest(t, ctx, store, "test-token"), + ProviderAuthSecretID: createMachinePoolProviderAuthSecretForTest( + t, + ctx, + store, + "test-token", + ), DefaultMachineMemoryMB: new(1024), MaxTotalMemoryMB: new(8192), MaxMachineMemoryMB: new(8192), } if provider == "blaxel" { input.ProviderConfig = json.RawMessage(`{"workspace":"test"}`) - input.DefaultMachineProviderOptions = json.RawMessage(`{"image":"test-image","region":"us-pdx-1"}`) + input.DefaultMachineProviderOptions = json.RawMessage( + `{"image":"test-image","region":"us-pdx-1"}`, + ) } else { input.DefaultMachineCPU = new(2) input.MaxTotalCPU = new(8) @@ -2428,11 +2506,17 @@ func TestCreatePoolMachineRejectsUnsupportedSizing(t *testing.T) { } machinePool, err := store.Execution().CreateMachinePool(ctx, input) require.NoError(t, err) - _, err = store.Execution().CreateProjectMachinePoolGrant(ctx, executionstore.CreateProjectMachinePoolGrantInput{ - OrgID: testOrgID, ProjectID: testProjectID, MachinePoolID: machinePool.ID, IdempotencyKey: "sizing", - }) + _, err = store.Execution(). + CreateProjectMachinePoolGrant(ctx, executionstore.CreateProjectMachinePoolGrantInput{ + OrgID: testOrgID, ProjectID: testProjectID, MachinePoolID: machinePool.ID, IdempotencyKey: "sizing", + }) require.NoError(t, err) - config := mustCreateAgentConfigFromYAML(t, ctx, store, agentPoolMachineConfigYAML(machinePool.Name, 4)) + config := mustCreateAgentConfigFromYAML( + t, + ctx, + store, + agentPoolMachineConfigYAML(machinePool.Name, 4), + ) agent, err := store.Execution().CreateAgentFixture(ctx, executionstore.AgentFixtureInput{ ProjectID: testProjectID, CurrentConfigID: config.ID, }) diff --git a/internal/storage/executionstore/process_action_completion_store.go b/internal/storage/executionstore/process_action_completion_store.go index 379f4652c..6906ca58d 100644 --- a/internal/storage/executionstore/process_action_completion_store.go +++ b/internal/storage/executionstore/process_action_completion_store.go @@ -9,8 +9,9 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/daemonprotocol" - "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/processresult" "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -45,7 +46,10 @@ func replayDaemonProcessActionReportTx( return DaemonProcessActionReportApplication{}, false, nil } if err != nil { - return DaemonProcessActionReportApplication{}, false, fmt.Errorf("load process action report replay: %w", err) + return DaemonProcessActionReportApplication{}, false, fmt.Errorf( + "load process action report replay: %w", + err, + ) } record := processActionRecordFromSQLC(row) if record.ID != input.ID || !isProcessActionTerminal(record.State) { @@ -62,7 +66,8 @@ func replayDaemonProcessActionReportTx( reportMatchesAction = false } case ProcessActionStateFailed, ProcessActionStateUnknown: - if record.StateReasonCode != input.StateReasonCode || record.StateReasonMessage != input.StateReasonMessage { + if record.StateReasonCode != input.StateReasonCode || + record.StateReasonMessage != input.StateReasonMessage { reportMatchesAction = false } outcome = ToolResultOutcomeFailed @@ -96,9 +101,9 @@ func replayDaemonProcessActionReportTx( processErr, ) } - result, err = canonicalProcessReadFailureResult( - processRecordFromSQLC(processRow), - record, + result, err = processresult.ReadFailure( + processRecordFromSQLC(processRow).resultFacts(), + record.resultFacts(), input.StateReasonCode, input.StateReasonMessage, ) @@ -109,11 +114,9 @@ func replayDaemonProcessActionReportTx( } } if len(result) == 0 || string(result) == "null" { - result, err = processActionToolResult( + result, err = processresult.ActionResult( input.ProcessID, - input.ID, - state, - input.StateReasonCode, + input.ID, string(state), input.StateReasonCode, errText, ) if err != nil { @@ -220,15 +223,19 @@ func (s *Store) completeDaemonProcessAction( if err := validateDaemonRuntimeAuthority(input.Authority); err != nil { return DaemonProcessActionReportApplication{}, err } - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ProcessID == uuid.Nil || input.ID == uuid.Nil { - return DaemonProcessActionReportApplication{}, errors.New("project, agent, process, and action are required") + if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ProcessID == uuid.Nil || + input.ID == uuid.Nil { + return DaemonProcessActionReportApplication{}, errors.New( + "project, agent, process, and action are required", + ) } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonProcessActionReportApplication{}, fmt.Errorf("begin complete process action: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { return DaemonProcessActionReportApplication{}, err @@ -242,11 +249,11 @@ func (s *Store) completeDaemonProcessAction( StateReasonMessage: input.StateReasonMessage, Result: input.Result, } - application, err := completeDaemonProcessActionTx(ctx, txNotifications, tx, qtx, completion, state) + application, err := completeDaemonProcessActionTx(ctx, unit, qtx, completion, state) if err != nil { return DaemonProcessActionReportApplication{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "complete process action"); err != nil { + if err := unit.Commit(ctx, "complete process action"); err != nil { return DaemonProcessActionReportApplication{}, err } return application, nil @@ -254,21 +261,20 @@ func (s *Store) completeDaemonProcessAction( func completeDaemonProcessActionTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, input processActionCompletionInput, state ProcessActionState, ) (DaemonProcessActionReportApplication, error) { - _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, - ID: input.AgentID, - }, - ) + _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: input.ProjectID, + ID: input.AgentID, + }, agentexecution.ExternalAuthority{}) if err != nil { - return DaemonProcessActionReportApplication{}, fmt.Errorf("lock agent for process action completion: %w", err) + return DaemonProcessActionReportApplication{}, fmt.Errorf( + "lock agent for process action completion: %w", + err, + ) } actionRow, err := qtx.GetProcessActionForReport( ctx, @@ -322,9 +328,9 @@ func completeDaemonProcessActionTx( process := processRecordFromSQLC(processRow) action := processActionRecordFromSQLC(actionRow) if state == ProcessActionStateApplied { - canonical, nextCursor, canonicalErr := canonicalProcessReadResult( - process, - action, + canonical, nextCursor, canonicalErr := processresult.Read( + process.resultFacts(), + action.resultFacts(), input.Result, ) if canonicalErr != nil { @@ -348,9 +354,9 @@ func completeDaemonProcessActionTx( reportAccepted = false } if state != ProcessActionStateApplied { - input.Result, err = canonicalProcessReadFailureResult( - process, - action, + input.Result, err = processresult.ReadFailure( + process.resultFacts(), + action.resultFacts(), input.StateReasonCode, input.StateReasonMessage, ) @@ -396,7 +402,10 @@ func completeDaemonProcessActionTx( }, ) default: - return DaemonProcessActionReportApplication{}, fmt.Errorf("unsupported process action completion state %q", state) + return DaemonProcessActionReportApplication{}, fmt.Errorf( + "unsupported process action completion state %q", + state, + ) } if errors.Is(err, pgx.ErrNoRows) { application, resolved, replayErr := replayDaemonProcessActionReportTx(ctx, qtx, input, state) @@ -442,121 +451,20 @@ func completeDaemonProcessActionTx( record := processActionRecordFromSQLC(row) resultCommitted := record.ToolCallID == uuid.Nil if record.ToolCallID != uuid.Nil { - outcome := ToolResultOutcomeSucceeded - errText := "" - if state != ProcessActionStateApplied { - outcome = ToolResultOutcomeFailed - errText = input.StateReasonMessage - if errText == "" { - errText = input.StateReasonCode - } - } - result := input.Result - if len(result) == 0 || string(result) == "null" { - result, err = processActionToolResult( - input.ProcessID, - input.ID, - state, - input.StateReasonCode, - errText, - ) - if err != nil { - return DaemonProcessActionReportApplication{}, err - } - } - contentParts, err := ToolResultContentParts(result) + h, err := unit.Handle(input.ProjectID, input.AgentID) if err != nil { return DaemonProcessActionReportApplication{}, err } - completedToolCall := false - toolRow, err := qtx.CompleteToolCallFromProcessAction( + completed, err := h.CompleteProcessAction( ctx, - dbsqlc.CompleteToolCallFromProcessActionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ToolCallID: record.ToolCallID, - ProcessID: input.ProcessID, - ProcessActionID: input.ID, - Outcome: string(outcome), - }, + agentexecution.ActionResult{ID: record.ID, Observed: input.Result}, ) if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - publication, checkErr := inspectPublishedToolCallResultTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - record.ToolCallID, - outcome, - contentParts, - ) - if checkErr != nil { - return DaemonProcessActionReportApplication{}, checkErr - } - if !publication.Published { - return DaemonProcessActionReportApplication{}, fmt.Errorf( - "terminal process action %s has no durable result for linked tool call %s", - record.ID, - record.ToolCallID, - ) - } - resultCommitted = publication.Matches - } else { - return DaemonProcessActionReportApplication{}, fmt.Errorf( - "complete tool call from process action: %w", - err, - ) - } - } else { - resultCommitted = true - completedToolCall = true - resultRecord := toolCallRecordFromProcessActionCompleteSQLC(toolRow) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx( - ctx, - txNotifications, - tx, - resultRecord, - nil, - ); err != nil { - return DaemonProcessActionReportApplication{}, err - } - } - if completedToolCall { - metadata, err := marshalJSON(struct { - Reason string `json:"reason"` - ProcessID uuid.UUID `json:"process_id"` - ProcessActionID uuid.UUID `json:"process_action_id"` - ToolCallID uuid.UUID `json:"tool_call_id"` - }{ - Reason: "process_action_result", - ProcessID: input.ProcessID, - ProcessActionID: input.ID, - ToolCallID: record.ToolCallID, - }) - if err != nil { - return DaemonProcessActionReportApplication{}, fmt.Errorf( - "marshal process action wakeup metadata: %w", - err, - ) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: metadata, - }, - ); err != nil { - return DaemonProcessActionReportApplication{}, fmt.Errorf( - "mark process action tool result wakeup: %w", - err, - ) - } + return DaemonProcessActionReportApplication{}, err } + resultCommitted = completed.Matches } - if resultCommitted && + if record.ToolCallID == uuid.Nil && resultCommitted && readNextCursor != nil && state == ProcessActionStateApplied { if _, err := qtx.AdvanceProcessDefaultOutputCursor( @@ -580,32 +488,6 @@ func completeDaemonProcessActionTx( }, nil } -func processActionToolResult( - processID uuid.UUID, - actionID uuid.UUID, - state ProcessActionState, - reasonCode string, - errText string, -) (json.RawMessage, error) { - body, err := marshalJSON(struct { - ProcessID string `json:"process_id"` - ProcessActionID string `json:"process_action_id"` - State ProcessActionState `json:"state"` - StateReasonCode string `json:"state_reason_code"` - Error string `json:"error"` - }{ - ProcessID: publicResourceID(publicid.KindProcess, processID), - ProcessActionID: publicResourceID(publicid.KindProcessAction, actionID), - State: state, - StateReasonCode: reasonCode, - Error: errText, - }) - if err != nil { - return nil, fmt.Errorf("marshal process action tool result: %w", err) - } - return body, nil -} - func publicResourceID(kind publicid.Kind, id uuid.UUID) string { encoded, err := publicid.Encode(kind, id) if err != nil { diff --git a/internal/storage/executionstore/process_actions_store.go b/internal/storage/executionstore/process_actions_store.go index afe8c067d..668b42428 100644 --- a/internal/storage/executionstore/process_actions_store.go +++ b/internal/storage/executionstore/process_actions_store.go @@ -12,7 +12,10 @@ import ( "github.com/omnara-ai/omnara/internal/storage/storeerr" ) -var ErrProcessMachineDeleted = storeerr.Tag(storeerr.ErrProcessTerminal, errors.New("process machine was deleted")) +var ErrProcessMachineDeleted = storeerr.Tag( + storeerr.ErrProcessTerminal, + errors.New("process machine was deleted"), +) func (t *toolCallTransaction) createProcessAction( ctx context.Context, @@ -45,6 +48,7 @@ func (t *toolCallTransaction) createProcessAction( return ProcessActionRecord{}, storeerr.ErrIdempotencyConflict } t.hasDurableCompletionOwner = true + t.owner.ActionID = record.ID if err := t.lockOrAcceptExisting(ctx); err != nil { return ProcessActionRecord{}, err } @@ -118,6 +122,7 @@ func (t *toolCallTransaction) createProcessAction( } record := processActionRecordFromSQLC(row) t.hasDurableCompletionOwner = true + t.owner.ActionID = record.ID t.requiresWaitingDisposition = true t.notifications.AddDaemonWork(processRow.MachineID) return record, nil @@ -171,7 +176,10 @@ func (t *toolCallTransaction) processActionCreateBlocker( return storeerr.ErrRuntimeLockInactive } -func (s *Store) ListDaemonProcessOffers(ctx context.Context, input DaemonWorkInput) ([]DaemonProcessOffer, error) { +func (s *Store) ListDaemonProcessOffers( + ctx context.Context, + input DaemonWorkInput, +) ([]DaemonProcessOffer, error) { if err := validateDaemonRuntimeAuthority(input.Authority); err != nil { return nil, err } @@ -286,7 +294,10 @@ func (s *Store) AcceptDaemonProcess( ); err != nil { if errors.Is(err, pgx.ErrNoRows) { if err := tx.Commit(ctx); err != nil { - return DaemonProcessOffer{}, false, fmt.Errorf("commit missing process for daemon accept: %w", err) + return DaemonProcessOffer{}, false, fmt.Errorf( + "commit missing process for daemon accept: %w", + err, + ) } return DaemonProcessOffer{}, false, nil } @@ -411,7 +422,10 @@ func (s *Store) AcceptDaemonProcessAction( } return DaemonProcessActionGrant{}, false, nil } - return DaemonProcessActionGrant{}, false, fmt.Errorf("lock machine for accept daemon process action: %w", err) + return DaemonProcessActionGrant{}, false, fmt.Errorf( + "lock machine for accept daemon process action: %w", + err, + ) } if _, err := qtx.LockDaemonProcessForAccept( ctx, @@ -464,7 +478,10 @@ func (s *Store) AcceptDaemonProcessAction( ) if errors.Is(err, pgx.ErrNoRows) { if err := tx.Commit(ctx); err != nil { - return DaemonProcessActionGrant{}, false, fmt.Errorf("commit accept daemon process action: %w", err) + return DaemonProcessActionGrant{}, false, fmt.Errorf( + "commit accept daemon process action: %w", + err, + ) } return DaemonProcessActionGrant{}, false, nil } @@ -519,7 +536,11 @@ func (s *Store) GetProcessActionByToolCall( ) (ProcessActionRecord, bool, error) { row, err := s.q.GetProcessActionByToolCall( ctx, - dbsqlc.GetProcessActionByToolCallParams{ProjectID: projectID, AgentID: agentID, ToolCallID: toolCallID}, + dbsqlc.GetProcessActionByToolCallParams{ + ProjectID: projectID, + AgentID: agentID, + ToolCallID: toolCallID, + }, ) if errors.Is(err, pgx.ErrNoRows) { return ProcessActionRecord{}, false, nil diff --git a/internal/storage/executionstore/process_cancellation_archive_integration_test.go b/internal/storage/executionstore/process_cancellation_archive_integration_test.go index 6503480e0..b7f9e0443 100644 --- a/internal/storage/executionstore/process_cancellation_archive_integration_test.go +++ b/internal/storage/executionstore/process_cancellation_archive_integration_test.go @@ -19,7 +19,14 @@ func TestCancelAgentOrdersCancelBeforeCanceledToolResult(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_orders_cause_before_effect") - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "cancel_orders_cause_before_effect", "run_command") + + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "cancel_orders_cause_before_effect", + "run_command", + ) _, err := startProcessForTest(ctx, fixture.Store, executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, @@ -143,17 +150,19 @@ func TestCancelAgentNoOpsTerminalTurnWithUncanceledRuntime(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_terminal_turn_live_runtime") - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"finish"}]`), - IdempotencyKey: "cancel-terminal-turn-live-runtime", - }) + + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"finish"}]`), + IdempotencyKey: "cancel-terminal-turn-live-runtime", + }) if err != nil { t.Fatalf("create content input: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, fixture.Store, @@ -164,15 +173,17 @@ func TestCancelAgentNoOpsTerminalTurnWithUncanceledRuntime(t *testing.T) { if !found { t.Fatal("expected admitted content input") } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent: %v", err) } - modelClaim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, OpeningInputIDs: []uuid.UUID{input.ID}, AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, OpeningInputIDs: []uuid.UUID{input.ID}, + }) + modelClaim := prepared1.Claim + if err != nil { t.Fatalf("claim model context: %v", err) } @@ -243,17 +254,19 @@ func TestCancelAgentWinsAgainstActiveModelCallAndRejectsLateAcceptance(t *testin t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_active_model_call") - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"cancel active model call"}]`), - IdempotencyKey: "cancel-active-model-call", - }) + + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"cancel active model call"}]`), + IdempotencyKey: "cancel-active-model-call", + }) if err != nil { t.Fatalf("create content input: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, fixture.Store, @@ -264,15 +277,17 @@ func TestCancelAgentWinsAgainstActiveModelCallAndRejectsLateAcceptance(t *testin if !found { t.Fatal("expected admitted content input") } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent: %v", err) } - modelClaim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, OpeningInputIDs: []uuid.UUID{input.ID}, AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) + prepared2, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, OpeningInputIDs: []uuid.UUID{input.ID}, + }) + modelClaim := prepared2.Claim + if err != nil { t.Fatalf("claim model context: %v", err) } @@ -364,6 +379,7 @@ func TestArchiveAgentAtomicallyStopsDurableModelCallWork(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "archive_model_call_"+test.name) + now := fixture.Now.Add(time.Minute) input, _, _, err := fixture.Store.Execution().CreateAgentContentInput( ctx, @@ -378,7 +394,7 @@ func TestArchiveAgentAtomicallyStopsDurableModelCallWork(t *testing.T) { if err != nil { t.Fatalf("create content input: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, fixture.Store, @@ -389,18 +405,19 @@ func TestArchiveAgentAtomicallyStopsDurableModelCallWork(t *testing.T) { if !found { t.Fatal("expected admitted content input") } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent: %v", err) } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) + prepared3, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, + }) + claim := prepared3.Claim + if err != nil { t.Fatalf("claim model context: %v", err) } @@ -462,14 +479,23 @@ func TestArchiveAgentAtomicallyStopsDurableModelCallWork(t *testing.T) { } if archived.ArchivedAt == nil || contextRecord.CompletedAt == nil || contextRecord.CompletedAt.Before(contextRecord.CreatedAt) { - t.Fatalf("archived model context has invalid timestamps: agent=%+v context=%+v", archived, contextRecord) + t.Fatalf( + "archived model context has invalid timestamps: agent=%+v context=%+v", + archived, + contextRecord, + ) } if test.recordRetry && archived.ArchivedAt.Before(*contextRecord.CompletedAt) { - t.Fatalf("agent archive predates the existing retry outcome: agent=%+v context=%+v", archived, contextRecord) + t.Fatalf( + "agent archive predates the existing retry outcome: agent=%+v context=%+v", + archived, + contextRecord, + ) } if !test.recordRetry && contextRecord.CompletedAt.Before(*archived.ArchivedAt) { t.Fatalf( - "archive-triggered model cancellation predates the agent archive: agent=%+v context=%+v", archived, + "archive-triggered model cancellation predates the agent archive: agent=%+v context=%+v", + archived, contextRecord, ) } @@ -501,6 +527,7 @@ func TestCancelAgentCancelsSteeringButPreservesQueuedBacklogWhenActive(t *testin t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_active_cancels_steering") + toolCallID := createToolCallForProcessTest( t, ctx, @@ -587,6 +614,7 @@ func TestCancelAgentStopsUnstartedTurnWithoutLiveRuntime(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_unstarted_turn_without_runtime") + if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( ctx, testProjectID, @@ -595,13 +623,14 @@ func TestCancelAgentStopsUnstartedTurnWithoutLiveRuntime(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"cancel me"}]`), - IdempotencyKey: "cancel-unstarted-turn", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"cancel me"}]`), + IdempotencyKey: "cancel-unstarted-turn", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -609,7 +638,8 @@ func TestCancelAgentStopsUnstartedTurnWithoutLiveRuntime(t *testing.T) { if err != nil { t.Fatalf("claim input work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) { t.Fatalf("claim = %+v found=%v, want executable input %s", claim, found, input.ID) } if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( @@ -668,7 +698,11 @@ func TestCancelAgentStopsUnstartedTurnWithoutLiveRuntime(t *testing.T) { if wakeups := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); wakeups != 0 { t.Fatalf("wakeups after canceling unstarted turn = %d, want 0", wakeups) } - if seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID); err != nil { + if seed, + found, + err := fixture.Store.Execution().NextAgentModelWork(ctx, + testProjectID, + fixture.AgentID); err != nil { t.Fatalf("load continuation seed after cancel: %v", err) } else if found { t.Fatalf("continuation seed after cancel = %+v, want none", seed) @@ -679,6 +713,7 @@ func TestCancelAgentReportsAlreadyPendingRuntimeCancel(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_already_canceled_runtime") + if _, err := requestAgentRuntimeCancelForTest( ctx, fixture.Store, @@ -726,6 +761,7 @@ func TestCancelAgentCanCancelLaterFrontierAfterPriorCancel(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "cancel_later_frontier_after_prior_cancel") + firstToolCallID := createToolCallForProcessTest( t, ctx, @@ -768,7 +804,8 @@ func TestCancelAgentCanCancelLaterFrontierAfterPriorCancel(t *testing.T) { if err != nil { t.Fatalf("first cancel: %v", err) } - firstToolCall, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, firstToolCallID) + firstToolCall, err := fixture.Store.Execution(). + GetToolCall(ctx, testProjectID, fixture.AgentID, firstToolCallID) if err != nil { t.Fatalf("load first tool call: %v", err) } @@ -843,7 +880,8 @@ func TestCancelAgentCanCancelLaterFrontierAfterPriorCancel(t *testing.T) { firstCancelResult.Event, ) } - secondToolCall, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, secondToolCallID) + secondToolCall, err := fixture.Store.Execution(). + GetToolCall(ctx, testProjectID, fixture.AgentID, secondToolCallID) if err != nil { t.Fatalf("load second tool call: %v", err) } diff --git a/internal/storage/executionstore/process_compaction_tool_completion_integration_test.go b/internal/storage/executionstore/process_compaction_tool_completion_integration_test.go index c288b63f3..ed52a2c22 100644 --- a/internal/storage/executionstore/process_compaction_tool_completion_integration_test.go +++ b/internal/storage/executionstore/process_compaction_tool_completion_integration_test.go @@ -48,18 +48,19 @@ func TestReplaceCompactionSourceRejectsLargerRange(t *testing.T) { if !found || len(admitted.Events) != 2 { t.Fatalf("admitted opening events = %+v found=%v, want two", admitted.Events, found) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent: %v", err) } - parent, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{inputs[0].ID, inputs[1].ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[1].Sequence, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{inputs[0].ID, inputs[1].ID}, + }) + parent := prepared1.Claim + if err != nil { t.Fatalf("claim blocked normal context: %v", err) } @@ -84,16 +85,17 @@ func TestReplaceCompactionSourceRejectsLargerRange(t *testing.T) { t.Fatalf("record compaction recovery: %v", err) } compactionClaim := handoff.CompactionCall - _, err = fixture.Store.Execution().ReplaceCompactionSource(ctx, executionstore.ReplaceCompactionSourceInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - ModelCallContextID: compactionClaim.Context.ID, - ErrorKind: "context_window", - ErrorCode: "candidate_projection_over_budget", - ErrorMessage: "candidate checkpoint did not fit", - NextSourceEventSequenceEnd: admitted.Events[1].Sequence, - }) + _, err = fixture.Store.Execution(). + ReplaceCompactionSource(ctx, executionstore.ReplaceCompactionSourceInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + ModelCallContextID: compactionClaim.Context.ID, + ErrorKind: "context_window", + ErrorCode: "candidate_projection_over_budget", + ErrorMessage: "candidate checkpoint did not fit", + NextSourceEventSequenceEnd: admitted.Events[1].Sequence, + }) if !errors.Is(err, storeerr.ErrStateTransitionConflict) { t.Fatalf("larger source replacement err = %v, want %v", err, storeerr.ErrStateTransitionConflict) } @@ -149,14 +151,15 @@ func TestSucceededCompactionCannotStartAnotherBranch(t *testing.T) { t.Fatalf("load compaction parent: found=%v err=%v", found, err) } - _, err = fixture.Store.Execution().ClaimCompactionModelCall(ctx, executionstore.ClaimCompactionModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - InputEventSequence: claim.Context.InputEventSequence, - SourceEventSequenceEnd: *claim.Context.SourceEventSequenceEnd - 1, - ParentContextID: parent.ID, - }) + _, err = fixture.Store.Execution(). + ClaimCompactionModelCall(ctx, executionstore.ClaimCompactionModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + InputEventSequence: claim.Context.InputEventSequence, + SourceEventSequenceEnd: *claim.Context.SourceEventSequenceEnd - 1, + ParentContextID: parent.ID, + }) if !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { t.Fatalf("second compaction branch error = %v, want %v", err, storeerr.ErrAgentNotAdvanceable) } @@ -195,12 +198,13 @@ func TestCompactionSourceAdjustmentFollowsRetryHistory(t *testing.T) { ); err != nil { t.Fatalf("record compaction retry: %v", err) } - retry, err := fixture.Store.Execution().ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - PredecessorModelCallContextID: first.Context.ID, - RuntimeLockID: fixture.Lock.ID, - }) + retry, err := fixture.Store.Execution(). + ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + PredecessorModelCallContextID: first.Context.ID, + RuntimeLockID: fixture.Lock.ID, + }) if err != nil { t.Fatalf("claim compaction retry: %v", err) } @@ -209,16 +213,17 @@ func TestCompactionSourceAdjustmentFollowsRetryHistory(t *testing.T) { } replacementEnd := *retry.Context.SourceEventSequenceEnd - 1 - replacement, err := fixture.Store.Execution().ReplaceCompactionSource(ctx, executionstore.ReplaceCompactionSourceInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - ModelCallContextID: retry.Context.ID, - ErrorKind: "context_window", - ErrorCode: "candidate_projection_over_budget", - ErrorMessage: "candidate checkpoint did not fit", - NextSourceEventSequenceEnd: replacementEnd, - }) + replacement, err := fixture.Store.Execution(). + ReplaceCompactionSource(ctx, executionstore.ReplaceCompactionSourceInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + ModelCallContextID: retry.Context.ID, + ErrorKind: "context_window", + ErrorCode: "candidate_projection_over_budget", + ErrorMessage: "candidate checkpoint did not fit", + NextSourceEventSequenceEnd: replacementEnd, + }) if err != nil || replacement.BoundaryPreempted { t.Fatalf("replace compaction source after retry: result=%+v err=%v", replacement, err) } diff --git a/internal/storage/executionstore/process_daemon_completion_store.go b/internal/storage/executionstore/process_daemon_completion_store.go index 27688cc8b..f1f761840 100644 --- a/internal/storage/executionstore/process_daemon_completion_store.go +++ b/internal/storage/executionstore/process_daemon_completion_store.go @@ -10,6 +10,8 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/daemonprotocol" + "github.com/omnara-ai/omnara/internal/processresult" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -31,25 +33,25 @@ func (s *Store) MarkProcessStarted( ) } input.SourceStartedAt = canonicalSourceTime(input.SourceStartedAt) - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonProcessReportApplication{}, fmt.Errorf("begin mark process started: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() qtx := dbsqlc.New(tx) if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { return DaemonProcessReportApplication{}, err } - _, err = qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, - ID: input.AgentID, - }, - ) + _, err = unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: input.ProjectID, + ID: input.AgentID, + }, agentexecution.ExternalAuthority{}) if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("lock agent for process start observation: %w", err) + return DaemonProcessReportApplication{}, fmt.Errorf( + "lock agent for process start observation: %w", + err, + ) } row, err := qtx.MarkProcessStarted( ctx, @@ -84,106 +86,29 @@ func (s *Store) MarkProcessStarted( record = processRecordFromStartedSQLC(row) } resultCommitted := false - var committedResult json.RawMessage if record.ToolCallID != uuid.Nil && record.ExecutionSpec.FileTransfer == nil { - startedRecord := record - startedRecord.State = ProcessStateRunning - result, err := startedProcessToolResult(startedRecord, input.Result) + h, err := unit.Handle(input.ProjectID, input.AgentID) if err != nil { return DaemonProcessReportApplication{}, err } - contentParts, err := ToolResultContentParts(result) + completed, err := h.CompleteProcess( + ctx, + agentexecution.ProcessResult{ID: record.ID, Started: true, Observed: input.Result}, + ) if err != nil { return DaemonProcessReportApplication{}, err } - outcome := ToolResultOutcomeSucceeded - toolRow, err := qtx.CompleteToolCallFromStartedProcess( + resultCommitted = completed.Matches + row, err := qtx.GetProcess( ctx, - dbsqlc.CompleteToolCallFromStartedProcessParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: record.ToolCallID, - ProcessID: record.ID, - Outcome: string(outcome), - }, + dbsqlc.GetProcessParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: record.ID}, ) if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - publication, checkErr := inspectPublishedToolCallResultTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - record.ToolCallID, - outcome, - contentParts, - ) - if checkErr != nil { - return DaemonProcessReportApplication{}, checkErr - } - if publication.Matches { - resultCommitted = true - committedResult = result - } else if !publication.Published { - return DaemonProcessReportApplication{}, fmt.Errorf( - "linked tool call %s has no durable result for started process %s", - record.ToolCallID, - record.ID, - ) - } - } else { - return DaemonProcessReportApplication{}, fmt.Errorf("complete tool call from started process: %w", err) - } - } else { - resultCommitted = true - committedResult = result - resultRecord := toolCallRecordFromStartedProcessCompleteSQLC(toolRow) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx( - ctx, - txNotifications, - tx, - resultRecord, - nil, - ); err != nil { - return DaemonProcessReportApplication{}, err - } - metadata, err := marshalJSON( - map[string]any{ - "reason": "process_started", - "process_id": record.ID, - "tool_call_id": record.ToolCallID, - }, - ) - if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("marshal process started wakeup metadata: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: metadata, - }, - ); err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("mark started process tool result wakeup: %w", err) - } - } - } - if resultCommitted { - if err := advanceProcessDefaultOutputCursorFromResult( - ctx, - qtx, - &record, - committedResult, - ); err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf( - "advance started process output cursor: %w", - err, - ) + return DaemonProcessReportApplication{}, err } + record = processRecordFromSQLC(row) } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "mark process started"); err != nil { + if err := unit.Commit(ctx, "mark process started"); err != nil { return DaemonProcessReportApplication{}, err } return DaemonProcessReportApplication{ @@ -222,25 +147,26 @@ func (s *Store) CompleteDaemonProcess( if !input.SourceEndedAt.IsZero() { input.SourceEndedAt = canonicalSourceTime(input.SourceEndedAt) } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DaemonProcessReportApplication{}, fmt.Errorf("begin complete daemon process: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { return DaemonProcessReportApplication{}, err } - _, err = qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, - ID: input.AgentID, - }, - ) + _, err = unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: input.ProjectID, + ID: input.AgentID, + }, agentexecution.ExternalAuthority{}) if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("lock agent for daemon process completion: %w", err) + return DaemonProcessReportApplication{}, fmt.Errorf( + "lock agent for daemon process completion: %w", + err, + ) } var startedAt *time.Time if !input.SourceStartedAt.IsZero() { @@ -292,93 +218,50 @@ func (s *Store) CompleteDaemonProcess( } reportMatchesProcess := processUpdated || daemonTerminalReportMatchesRecord(record, input) resultCommitted := false - var committedResult json.RawMessage if reportMatchesProcess && record.ToolCallID != uuid.Nil { - outcome, result, resultErr := processToolResult(record) - if resultErr != nil { - return DaemonProcessReportApplication{}, resultErr - } - if len(input.Result) > 0 && string(input.Result) != "null" { - result, err = commandTerminalToolResult(record.ID, input.Result) + request := agentexecution.ProcessResult{ID: record.ID, Observed: input.Result} + if record.ExecutionSpec.FileTransfer != nil { + _, result, err := processresult.Terminal(record.resultFacts()) if err != nil { return DaemonProcessReportApplication{}, err } + if len(input.Result) > 0 && string(input.Result) != "null" { + result, err = commandTerminalToolResult(record.ID, input.Result) + if err != nil { + return DaemonProcessReportApplication{}, err + } + } + outcome, raw, err := fileTransferToolResultContentParts(ctx, qtx, record, result) + if err != nil { + return DaemonProcessReportApplication{}, err + } + parts, err := parseToolResultContentBlocks(raw) + if err != nil { + return DaemonProcessReportApplication{}, err + } + request.Content, err = executionContent(parts) + if err != nil { + return DaemonProcessReportApplication{}, err + } + request.Outcome = string(outcome) } - var contentParts json.RawMessage - if record.ExecutionSpec.FileTransfer != nil { - outcome, contentParts, err = fileTransferToolResultContentParts(ctx, qtx, record, result) - } else { - contentParts, err = ToolResultContentParts(result) + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return DaemonProcessReportApplication{}, err } + completed, err := h.CompleteProcess(ctx, request) if err != nil { return DaemonProcessReportApplication{}, err } - toolRow, err := qtx.CompleteToolCallFromProcess( + resultCommitted = completed.Matches + row, err := qtx.GetProcess( ctx, - dbsqlc.CompleteToolCallFromProcessParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: record.ToolCallID, - ProcessID: record.ID, - Outcome: string(outcome), - }, + dbsqlc.GetProcessParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: record.ID}, ) if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - publication, checkErr := inspectPublishedToolCallResultTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - record.ToolCallID, - outcome, - contentParts, - ) - if checkErr != nil { - return DaemonProcessReportApplication{}, checkErr - } - if publication.Matches { - resultCommitted = true - committedResult = result - } else if !publication.Published { - return DaemonProcessReportApplication{}, fmt.Errorf( - "linked tool call %s has no durable result for process %s", - record.ToolCallID, - record.ID, - ) - } - } else { - return DaemonProcessReportApplication{}, fmt.Errorf("complete tool call from daemon process: %w", err) - } - } else { - resultCommitted = true - committedResult = result - resultRecord := toolCallRecordFromProcessCompleteSQLC(toolRow) - resultRecord.ResultContentParts = contentParts - if _, err := appendToolResultEventTx(ctx, txNotifications, tx, resultRecord, nil); err != nil { - return DaemonProcessReportApplication{}, err - } - metadata, err := marshalJSON( - map[string]any{ - "reason": "process_result", - "process_id": record.ID, - "tool_call_id": record.ToolCallID, - }, - ) - if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("marshal process result wakeup metadata: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: metadata, - }, - ); err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("mark process tool result wakeup: %w", err) - } + return DaemonProcessReportApplication{}, err } + record = processRecordFromSQLC(row) } else if record.ToolCallID != uuid.Nil { published, checkErr := publishedToolCallResultExistsTx( ctx, @@ -397,24 +280,10 @@ func (s *Store) CompleteDaemonProcess( ) } } - if resultCommitted { - if err := advanceProcessDefaultOutputCursorFromResult( - ctx, - qtx, - &record, - committedResult, - ); err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf( - "advance completed process output cursor: %w", - err, - ) - } - } if input.StorageExhausted { if err := completeUnresolvedProcessActionsForClosedProcessTx( ctx, - txNotifications, - tx, + unit, qtx, record.OrgID, record.ID, @@ -425,15 +294,14 @@ func (s *Store) CompleteDaemonProcess( } else if processUpdated { if err := completeQueuedProcessActionsForTerminalProcessTx( ctx, - txNotifications, - tx, + unit, qtx, record, ); err != nil { return DaemonProcessReportApplication{}, err } } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "complete daemon process"); err != nil { + if err := unit.Commit(ctx, "complete daemon process"); err != nil { return DaemonProcessReportApplication{}, err } return DaemonProcessReportApplication{ @@ -569,7 +437,11 @@ func publishedToolCallResultExistsTx( } result, err := qtx.GetToolCallResultByToolCall( ctx, - dbsqlc.GetToolCallResultByToolCallParams{ProjectID: projectID, AgentID: agentID, ToolCallID: toolCallID}, + dbsqlc.GetToolCallResultByToolCallParams{ + ProjectID: projectID, + AgentID: agentID, + ToolCallID: toolCallID, + }, ) if errors.Is(err, pgx.ErrNoRows) { return false, nil @@ -579,7 +451,11 @@ func publishedToolCallResultExistsTx( } eventExists, err := qtx.ToolCallResultHasTypedEvent( ctx, - dbsqlc.ToolCallResultHasTypedEventParams{ProjectID: projectID, AgentID: agentID, ToolCallResultID: &result.ID}, + dbsqlc.ToolCallResultHasTypedEventParams{ + ProjectID: projectID, + AgentID: agentID, + ToolCallResultID: &result.ID, + }, ) if err != nil { return false, fmt.Errorf("check linked tool result event: %w", err) @@ -650,13 +526,11 @@ func inspectPublishedToolCallResultTx( if !eventExists { return toolCallResultPublication{}, nil } - storedParts, err := toolCallResultContentBlocksTx( + tool, err := qtx.GetToolCall( ctx, - qtx, - projectID, - agentID, - result.ID, + dbsqlc.GetToolCallParams{ProjectID: projectID, AgentID: agentID, ID: toolCallID}, ) + storedParts := tool.ResultContentParts if err != nil { return toolCallResultPublication{}, err } @@ -666,114 +540,3 @@ func inspectPublishedToolCallResultTx( sameJSON(storedParts, normalizedJSON(contentParts)), }, nil } - -func completedToolCallMissIsBenignTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, toolCallID uuid.UUID, - allowStartedProcessResult bool, -) (bool, error) { - existing, err := qtx.GetToolCall( - ctx, - dbsqlc.GetToolCallParams{ProjectID: projectID, AgentID: agentID, ID: toolCallID}, - ) - if err != nil { - return false, fmt.Errorf("load linked tool call after completion miss: %w", err) - } - if existing.State != "completed" || existing.CompletedAt == nil { - return false, nil - } - result, err := qtx.GetToolCallResultByToolCall( - ctx, - dbsqlc.GetToolCallResultByToolCallParams{ - ProjectID: projectID, - AgentID: agentID, - ToolCallID: toolCallID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return false, nil - } - if err != nil { - return false, fmt.Errorf("load linked tool result after completion miss: %w", err) - } - contentParts, err := toolCallResultContentBlocksTx( - ctx, - qtx, - projectID, - agentID, - result.ID, - ) - if err != nil { - return false, err - } - canceled := result.Outcome == string(ToolResultOutcomeCanceled) - machineUnreachable := structuredContentStringFieldEquals( - contentParts, - "error_code", - ProcessToolReasonMachineUnreachable, - ) - startedProcess := allowStartedProcessResult && - result.Outcome == string(ToolResultOutcomeSucceeded) && - structuredContentStringFieldEquals( - contentParts, - "state", - string(ProcessStateRunning), - ) - if !canceled && !machineUnreachable && !startedProcess { - return false, nil - } - eventExists, err := qtx.ToolCallResultHasTypedEvent( - ctx, - dbsqlc.ToolCallResultHasTypedEventParams{ - ProjectID: projectID, - AgentID: agentID, - ToolCallResultID: &result.ID, - }, - ) - if err != nil { - return false, fmt.Errorf( - "check linked tool result event after completion miss: %w", - err, - ) - } - return eventExists, nil -} - -func structuredContentStringFieldEquals( - raw json.RawMessage, - field string, - want string, -) bool { - var parts []struct { - Type string `json:"type"` - Value json.RawMessage `json:"value"` - } - if json.Unmarshal(raw, &parts) != nil { - return false - } - for _, part := range parts { - if part.Type == "structured_data" && - jsonObjectStringFieldEquals(part.Value, field, want) { - return true - } - } - return false -} - -func jsonObjectStringFieldEquals( - raw json.RawMessage, - field string, - want string, -) bool { - var object map[string]json.RawMessage - if json.Unmarshal(raw, &object) != nil { - return false - } - value, ok := object[field] - if !ok { - return false - } - var got string - return json.Unmarshal(value, &got) == nil && got == want -} diff --git a/internal/storage/executionstore/process_model_ready_frontier_races_integration_test.go b/internal/storage/executionstore/process_model_ready_frontier_races_integration_test.go index 31e3fb8e1..88ff38f0c 100644 --- a/internal/storage/executionstore/process_model_ready_frontier_races_integration_test.go +++ b/internal/storage/executionstore/process_model_ready_frontier_races_integration_test.go @@ -16,6 +16,7 @@ func TestQueuedInputWaitsForAmbiguousCompactionRetry(t *testing.T) { t.Parallel() ctx := context.Background() fixture, compaction := sentCompactionFrontierFixture(t, ctx, "queued_waits_compaction") + queued := createFrontierRaceInput( t, ctx, @@ -48,7 +49,8 @@ func TestQueuedInputWaitsForAmbiguousCompactionRetry(t *testing.T) { if err != nil { t.Fatalf("claim ambiguous compaction retry: %v", err) } - if !found || work.Kind != executionstore.AgentWorkModel || work.Model.Kind != executionstore.ModelWorkResume || + if !found || work.Kind != executionstore.AgentWorkModel || + work.Model.Kind != executionstore.ModelWorkResume || work.Model.ModelCallContextID != compaction.Context.ID { t.Fatalf( "claimed work = %+v found=%v, want compaction context %s before queued input", @@ -100,7 +102,8 @@ func TestSteeringStartsNewFrontierAfterAmbiguousNormalSendWithCapacity(t *testin if err != nil { t.Fatalf("claim steering after ambiguous normal send: %v", err) } - if !found || work.Kind != executionstore.AgentWorkModel || work.Model.Kind != executionstore.ModelWorkStart || + if !found || work.Kind != executionstore.AgentWorkModel || + work.Model.Kind != executionstore.ModelWorkStart || work.Model.AdmittedInputTurn.Turn.ID == uuid.Nil || len(work.Model.InputIDs) == 0 || work.Model.InputIDs[len(work.Model.InputIDs)-1] != steering.ID { @@ -118,6 +121,7 @@ func TestSteeringStartsNewFrontierAfterAmbiguousCompactionSendWithCapacity(t *te t.Parallel() ctx := context.Background() fixture, compaction := sentCompactionFrontierFixture(t, ctx, "steering_ambiguous_compaction") + parent, found, err := fixture.Store.Execution().GetNormalModelCallContextForFrontier( ctx, testProjectID, @@ -159,14 +163,21 @@ func TestSteeringStartsNewFrontierAfterAmbiguousCompactionSendWithCapacity(t *te if err != nil { t.Fatalf("claim steering after ambiguous compaction send: %v", err) } - if !found || work.Kind != executionstore.AgentWorkModel || work.Model.Kind != executionstore.ModelWorkStart || + if !found || work.Kind != executionstore.AgentWorkModel || + work.Model.Kind != executionstore.ModelWorkStart || work.Model.AdmittedInputTurn.Turn.ID == uuid.Nil || len(work.Model.InputIDs) == 0 || work.Model.InputIDs[len(work.Model.InputIDs)-1] != steering.ID { t.Fatalf("steering work = %+v found=%v, want fresh admitted turn", work, found) } fresh := claimTestNormalModelCallForWork(t, ctx, fixture, work, time.Now().UTC()) - assertFrontierRaceContextState(t, ctx, fixture, compaction.Context.ID, executionstore.ModelCallContextFailed) + assertFrontierRaceContextState( + t, + ctx, + fixture, + compaction.Context.ID, + executionstore.ModelCallContextFailed, + ) assertFrontierRaceContextState( t, ctx, @@ -184,6 +195,7 @@ func TestCompletedCompactionPublicationWinsBeforeWaitingSteering(t *testing.T) { t.Parallel() ctx := context.Background() fixture, compaction := sentCompactionFrontierFixture(t, ctx, "completed_compaction_wins") + parent, found, err := fixture.Store.Execution().GetNormalModelCallContextForFrontier( ctx, testProjectID, @@ -216,7 +228,13 @@ func TestCompletedCompactionPublicationWinsBeforeWaitingSteering(t *testing.T) { if checkpoint.ID == uuid.Nil || checkpoint.CheckpointEventID == uuid.Nil { t.Fatalf("published checkpoint = %+v, want durable checkpoint and event", checkpoint) } - assertFrontierRaceContextState(t, ctx, fixture, compaction.Context.ID, executionstore.ModelCallContextSucceeded) + assertFrontierRaceContextState( + t, + ctx, + fixture, + compaction.Context.ID, + executionstore.ModelCallContextSucceeded, + ) assertFrontierRaceContextState( t, ctx, @@ -241,7 +259,8 @@ func TestCompletedCompactionPublicationWinsBeforeWaitingSteering(t *testing.T) { if err != nil { t.Fatalf("claim steering after accepted checkpoint: %v", err) } - if !found || work.Kind != executionstore.AgentWorkModel || work.Model.Kind != executionstore.ModelWorkStart || + if !found || work.Kind != executionstore.AgentWorkModel || + work.Model.Kind != executionstore.ModelWorkStart || work.Model.AdmittedInputTurn.Turn.ID == uuid.Nil || len(work.Model.InputIDs) == 0 || work.Model.InputIDs[len(work.Model.InputIDs)-1] != steering.ID { @@ -278,14 +297,15 @@ func createFrontierRaceInput( now time.Time, ) executionstore.AgentInputRecord { t.Helper() - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"new frontier input"}]`), - DeliveryMode: deliveryMode, - IdempotencyKey: idempotencyKey, - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"new frontier input"}]`), + DeliveryMode: deliveryMode, + IdempotencyKey: idempotencyKey, + }) if err != nil { t.Fatalf("create %s input: %v", deliveryMode, err) } diff --git a/internal/storage/executionstore/process_read_results.go b/internal/storage/executionstore/process_read_results.go index f317b52ae..8df93434a 100644 --- a/internal/storage/executionstore/process_read_results.go +++ b/internal/storage/executionstore/process_read_results.go @@ -3,158 +3,12 @@ package executionstore import ( "context" "encoding/json" - "errors" - "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/processaction" - "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/processresult" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) -type daemonProcessReadObservation struct { - ProcessID string `json:"process_id"` - Output *string `json:"output"` - Cursor *int64 `json:"cursor"` - NextCursor *int64 `json:"next_cursor"` - Truncated *bool `json:"truncated"` -} - -func canonicalProcessReadResult( - process ProcessRecord, - action ProcessActionRecord, - raw json.RawMessage, -) (json.RawMessage, *int64, error) { - payload, err := processaction.DecodeReadPayload(action.Payload) - if err != nil { - return nil, nil, fmt.Errorf("decode read request: %w", err) - } - - var observed daemonProcessReadObservation - if err := json.Unmarshal(raw, &observed); err != nil { - return nil, nil, fmt.Errorf("decode read observation: %w", err) - } - if observed.Output == nil || - observed.Cursor == nil || - observed.NextCursor == nil || - observed.Truncated == nil { - return nil, nil, errors.New("read observation is incomplete") - } - publicProcessID := publicResourceID(publicid.KindProcess, process.ID) - if observed.ProcessID != publicProcessID { - return nil, nil, errors.New("read observation names a different process") - } - minimumCursor := process.DefaultOutputCursor - if payload.Cursor != nil { - minimumCursor = *payload.Cursor - } - if *observed.Cursor < 0 || - *observed.Cursor < minimumCursor || - (*observed.Cursor > minimumCursor && !*observed.Truncated) || - *observed.NextCursor < *observed.Cursor { - return nil, nil, errors.New("read observation cursor range is invalid") - } - span := *observed.NextCursor - *observed.Cursor - maxBytes := payload.ObservationBytes() - if span > int64(maxBytes) || - len([]byte(*observed.Output)) > maxBytes || - int64(len([]byte(*observed.Output))) > span || - (span > 0 && *observed.Output == "") { - return nil, nil, errors.New("read observation exceeds its requested range") - } - - result := map[string]any{ - "process_id": publicProcessID, - "state": process.State, - "output": *observed.Output, - "cursor": *observed.Cursor, - "next_cursor": *observed.NextCursor, - "truncated": *observed.Truncated, - "done": isProcessTerminal(process.State), - } - if process.SourceStartedAt != nil { - result["started_at"] = *process.SourceStartedAt - } - if isProcessTerminal(process.State) { - result["ended_at"] = process.SourceEndedAt - result["exit_code"] = process.ExitCode - result["exit_signal"] = process.ExitSignal - result["state_reason_code"] = process.StateReasonCode - result["state_reason_message"] = process.StateReasonMessage - } - canonical, err := marshalJSON(result) - if err != nil { - return nil, nil, err - } - if payload.Cursor != nil { - return canonical, nil, nil - } - next := *observed.NextCursor - return canonical, &next, nil -} - -func canonicalProcessReadFailureResult( - process ProcessRecord, - action ProcessActionRecord, - code string, - message string, -) (json.RawMessage, error) { - result := map[string]any{ - "process_id": publicResourceID(publicid.KindProcess, process.ID), - "process_action_id": publicResourceID(publicid.KindProcessAction, action.ID), - "state": process.State, - "error_code": code, - "message": message, - "error": message, - "retryable": code == ProcessToolReasonMachineUnreachable, - "done": isProcessTerminal(process.State), - } - if process.SourceStartedAt != nil { - result["started_at"] = *process.SourceStartedAt - } - if isProcessTerminal(process.State) { - result["ended_at"] = process.SourceEndedAt - result["exit_code"] = process.ExitCode - result["exit_signal"] = process.ExitSignal - result["state_reason_code"] = process.StateReasonCode - result["state_reason_message"] = process.StateReasonMessage - } - return marshalJSON(result) -} - -func processReadObservationMatchesResult( - raw json.RawMessage, - result json.RawMessage, -) (bool, error) { - observed, ok := decodeProcessReadObservation(raw) - if !ok { - return false, nil - } - var committed daemonProcessReadObservation - if err := json.Unmarshal(result, &committed); err != nil { - return false, fmt.Errorf("decode committed read result: %w", err) - } - if observed.Output == nil || committed.Output == nil || - observed.Cursor == nil || committed.Cursor == nil || - observed.NextCursor == nil || committed.NextCursor == nil || - observed.Truncated == nil || committed.Truncated == nil { - return false, nil - } - return observed.ProcessID == committed.ProcessID && - *observed.Output == *committed.Output && - *observed.Cursor == *committed.Cursor && - *observed.NextCursor == *committed.NextCursor && - *observed.Truncated == *committed.Truncated, nil -} - -func decodeProcessReadObservation( - raw json.RawMessage, -) (daemonProcessReadObservation, bool) { - var observed daemonProcessReadObservation - return observed, json.Unmarshal(raw, &observed) == nil -} - func inspectPublishedProcessReadObservationTx( ctx context.Context, qtx *dbsqlc.Queries, @@ -163,30 +17,11 @@ func inspectPublishedProcessReadObservationTx( toolCallID uuid.UUID, raw json.RawMessage, ) (toolCallResultPublication, error) { - result, err := qtx.GetToolCallResultByToolCall( + tool, err := qtx.GetToolCall( ctx, - dbsqlc.GetToolCallResultByToolCallParams{ - ProjectID: projectID, - AgentID: agentID, - ToolCallID: toolCallID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return toolCallResultPublication{}, nil - } - if err != nil { - return toolCallResultPublication{}, fmt.Errorf( - "load committed read result: %w", - err, - ) - } - contentParts, err := toolCallResultContentBlocksTx( - ctx, - qtx, - projectID, - agentID, - result.ID, + dbsqlc.GetToolCallParams{ProjectID: projectID, AgentID: agentID, ID: toolCallID}, ) + contentParts := tool.ResultContentParts if err != nil { return toolCallResultPublication{}, err } @@ -207,7 +42,7 @@ func inspectPublishedProcessReadObservationTx( publication.Matches = false return publication, nil } - publication.Matches, err = processReadObservationMatchesResult( + publication.Matches, err = processresult.ReadMatches( raw, committed, ) @@ -229,80 +64,3 @@ func structuredToolResultValue(parts json.RawMessage) (json.RawMessage, bool) { } return nil, false } - -func processObservationNextCursor( - raw json.RawMessage, -) (int64, bool, error) { - if len(raw) == 0 || string(raw) == "null" || string(raw) == "{}" { - return 0, false, nil - } - var observed struct { - Output *string `json:"output"` - Cursor *int64 `json:"cursor"` - NextCursor *int64 `json:"next_cursor"` - } - if err := json.Unmarshal(raw, &observed); err != nil { - return 0, false, err - } - if observed.NextCursor == nil { - if observed.Cursor != nil || observed.Output != nil { - return 0, false, errors.New( - "process observation cursor range is incomplete", - ) - } - return 0, false, nil - } - if observed.Cursor == nil || observed.Output == nil || - *observed.Cursor < 0 || - *observed.NextCursor < *observed.Cursor { - return 0, false, errors.New( - "process observation cursor range is invalid", - ) - } - span := *observed.NextCursor - *observed.Cursor - outputBytes := len([]byte(*observed.Output)) - if span > int64(processaction.MaxObservationBytes) || - outputBytes > processaction.MaxObservationBytes || - int64(outputBytes) > span || - (span > 0 && *observed.Output == "") { - return 0, false, errors.New( - "process observation exceeds its allowed range", - ) - } - return *observed.NextCursor, true, nil -} - -func advanceProcessDefaultOutputCursorFromResult( - ctx context.Context, - qtx *dbsqlc.Queries, - process *ProcessRecord, - result json.RawMessage, -) error { - nextCursor, found, err := processObservationNextCursor(result) - if err != nil { - return fmt.Errorf("validate process output cursor: %w", err) - } - if !found { - return nil - } - updated, err := qtx.AdvanceProcessDefaultOutputCursor( - ctx, - dbsqlc.AdvanceProcessDefaultOutputCursorParams{ - ProjectID: process.ProjectID, - AgentID: process.AgentID, - ID: process.ID, - NextCursor: nextCursor, - }, - ) - if err != nil { - return err - } - if updated != 1 { - return errors.New("process disappeared while advancing its output cursor") - } - process.DefaultOutputCursor = max( - process.DefaultOutputCursor, - nextCursor, - ) - return nil -} diff --git a/internal/storage/executionstore/process_result_facts.go b/internal/storage/executionstore/process_result_facts.go new file mode 100644 index 000000000..5c8bdf279 --- /dev/null +++ b/internal/storage/executionstore/process_result_facts.go @@ -0,0 +1,21 @@ +package executionstore + +import "github.com/omnara-ai/omnara/internal/processresult" + +func (p ProcessRecord) resultFacts() processresult.Process { + return processresult.Process{ + ID: p.ID, + State: string(p.State), + DefaultOutputCursor: p.DefaultOutputCursor, + SourceStartedAt: p.SourceStartedAt, + SourceEndedAt: p.SourceEndedAt, + ExitCode: p.ExitCode, + ExitSignal: p.ExitSignal, + StateReasonCode: p.StateReasonCode, + StateReasonMessage: p.StateReasonMessage, + ExecutionSpec: p.ExecutionSpec, + } +} +func (a ProcessActionRecord) resultFacts() processresult.Action { + return processresult.Action{ID: a.ID, Payload: a.Payload} +} diff --git a/internal/storage/executionstore/process_retry_model_ready_integration_test.go b/internal/storage/executionstore/process_retry_model_ready_integration_test.go index ba6817b1c..dc31fbf63 100644 --- a/internal/storage/executionstore/process_retry_model_ready_integration_test.go +++ b/internal/storage/executionstore/process_retry_model_ready_integration_test.go @@ -20,19 +20,20 @@ import ( func TestConcurrentRetryClaimsReuseOneDurableContext(t *testing.T) { t.Parallel() ctx := context.Background() - fixture, admitted, agent := newMultiInputContinuationSeedFixture(t, ctx, "concurrent_retry_claim") + fixture, admitted, _ := newMultiInputContinuationSeedFixture(t, ctx, "concurrent_retry_claim") openingInputIDs := make([]uuid.UUID, 0, len(admitted.Inputs)) for _, input := range admitted.Inputs { openingInputIDs = append(openingInputIDs, input.ID) } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[len(admitted.Events)-1].Sequence, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, + }) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim initial model call: %v", err) } @@ -111,11 +112,19 @@ func TestConcurrentRetryClaimsReuseOneDurableContext(t *testing.T) { if first.claim.Context.ID != second.claim.Context.ID || first.claim.Context.ID == claim.Context.ID || first.claim.Context.AttemptNumber != 2 || second.claim.Context.AttemptNumber != 2 { - t.Fatalf("concurrent retry claims = %+v / %+v, want one shared retry context 2", first.claim, second.claim) + t.Fatalf( + "concurrent retry claims = %+v / %+v, want one shared retry context 2", + first.claim, + second.claim, + ) } if first.claim.Created == second.claim.Created || first.claim.Claimed == second.claim.Claimed || first.claim.Created != first.claim.Claimed || second.claim.Created != second.claim.Claimed { - t.Fatalf("concurrent retry claims = %+v / %+v, want exactly one creator and sender", first.claim, second.claim) + t.Fatalf( + "concurrent retry claims = %+v / %+v, want exactly one creator and sender", + first.claim, + second.claim, + ) } var contextRows, liveRows int if err := fixture.Store.pool.QueryRow(ctx, ` @@ -167,10 +176,10 @@ func TestRetryClaimRejectedAfterStopEvent(t *testing.T) { var turnID uuid.UUID if err := fixture.Store.pool.QueryRow(ctx, ` SELECT turn_id -FROM model_call_context_turns +FROM model_call_contexts WHERE project_id = $1 AND agent_id = $2 - AND model_call_context_id = $3 + AND id = $3 `, testProjectID, fixture.AgentID, claim.Context.ID).Scan(&turnID); err != nil { t.Fatalf("load failed context turn: %v", err) } @@ -182,12 +191,13 @@ WHERE project_id = $1 retryAt.Add(time.Millisecond), ) - next, err := fixture.Store.Execution().ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - PredecessorModelCallContextID: claim.Context.ID, - RuntimeLockID: fixture.Lock.ID, - }) + next, err := fixture.Store.Execution(). + ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + PredecessorModelCallContextID: claim.Context.ID, + RuntimeLockID: fixture.Lock.ID, + }) if err != nil { t.Fatalf("claim retry after stop: %v", err) } @@ -311,13 +321,14 @@ func TestParentRetryWaitsWhileCompactionContextIsLive(t *testing.T) { t.Fatalf("release one-live fixture runtime: %v", err) } now := time.Now().UTC() - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise runtime context uniqueness"}]`), - IdempotencyKey: "one-live-model-context-per-runtime", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise runtime context uniqueness"}]`), + IdempotencyKey: "one-live-model-context-per-runtime", + }) if err != nil { t.Fatalf("create one-live input: %v", err) } @@ -325,7 +336,8 @@ func TestParentRetryWaitsWhileCompactionContextIsLive(t *testing.T) { ctx, testClaimNextAgentWorkInput(), ) - if err != nil || !found || work.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(work, input.ID) { + if err != nil || !found || work.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(work, input.ID) { t.Fatalf("claim one-live work = %+v found=%v err=%v", work, found, err) } parent := claimTestNormalModelCallForWork(t, ctx, fixture, work, now.Add(2*time.Second)) @@ -353,12 +365,13 @@ func TestParentRetryWaitsWhileCompactionContextIsLive(t *testing.T) { t.Fatalf("one-live compaction handoff = %+v", handoff) } - retry, err := fixture.Store.Execution().ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - PredecessorModelCallContextID: parent.Context.ID, - RuntimeLockID: work.RuntimeLock.ID, - }) + retry, err := fixture.Store.Execution(). + ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + PredecessorModelCallContextID: parent.Context.ID, + RuntimeLockID: work.RuntimeLock.ID, + }) if err != nil { t.Fatalf("claim parent retry during live compaction: %v", err) } @@ -510,13 +523,14 @@ func TestTerminalModelCallErrorStopsUntilNewInput(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"try the provider"}]`), - IdempotencyKey: "terminal-model-error", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"try the provider"}]`), + IdempotencyKey: "terminal-model-error", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -524,7 +538,8 @@ func TestTerminalModelCallErrorStopsUntilNewInput(t *testing.T) { if err != nil { t.Fatalf("claim input work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) { t.Fatalf("claim = %+v found=%v, want executable input %s", claim, found, input.ID) } modelClaim := claimTestNormalModelCallForWork(t, ctx, fixture, claim, now.Add(2*time.Second)) @@ -543,14 +558,16 @@ func TestTerminalModelCallErrorStopsUntilNewInput(t *testing.T) { ErrorDetails: json.RawMessage(`{"retryable":false}`), ProviderReportedCostUSD: "0.0000015", } - errorEvent, err := fixture.Store.Execution().RecordModelCallErrorAndCompleteContext(ctx, terminalErrorInput) + errorEvent, err := fixture.Store.Execution(). + RecordModelCallErrorAndCompleteContext(ctx, terminalErrorInput) if err != nil { t.Fatalf("record terminal model error: %v", err) } if errorEvent.Kind != "model_output" { t.Fatalf("terminal error event = %+v, want model_output", errorEvent) } - replayedErrorEvent, err := fixture.Store.Execution().RecordModelCallErrorAndCompleteContext(ctx, terminalErrorInput) + replayedErrorEvent, err := fixture.Store.Execution(). + RecordModelCallErrorAndCompleteContext(ctx, terminalErrorInput) if err != nil { t.Fatalf("replay terminal model error after ambiguous commit: %v", err) } @@ -639,7 +656,11 @@ func TestTerminalModelCallErrorStopsUntilNewInput(t *testing.T) { if wakeups := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); wakeups != 0 { t.Fatalf("terminal-model-error wakeups = %d, want 0", wakeups) } - if seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID); err != nil { + if seed, + found, + err := fixture.Store.Execution().NextAgentModelWork(ctx, + testProjectID, + fixture.AgentID); err != nil { t.Fatalf("load seed after terminal model error: %v", err) } else if found { t.Fatalf("terminal model error continuation seed = %+v, want none", seed) @@ -664,7 +685,8 @@ func TestTerminalModelCallErrorStopsUntilNewInput(t *testing.T) { if err != nil { t.Fatalf("claim input after terminal model error: %v", err) } - if !found || newWork.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(newWork, newInput.ID) || + if !found || newWork.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(newWork, newInput.ID) || newWork.Model.TurnID == claim.Model.TurnID { t.Fatalf("new work after terminal model error = %+v found=%v, want a fresh turn", newWork, found) } diff --git a/internal/storage/executionstore/process_runtime_integration_test.go b/internal/storage/executionstore/process_runtime_integration_test.go index 161432978..39fb087e4 100644 --- a/internal/storage/executionstore/process_runtime_integration_test.go +++ b/internal/storage/executionstore/process_runtime_integration_test.go @@ -225,7 +225,13 @@ WHERE org_id = $1 AND machine_id = $2 AND id = $3 ) done <- acceptResult{found: found, err: acceptErr} }() - integrationdb.WaitForLockWaitBlockedBy(t, ctx, fixture.Store.pool, "-- name: LockDaemonProcessForAccept", blockingPID) + integrationdb.WaitForLockWaitBlockedBy( + t, + ctx, + fixture.Store.pool, + "-- name: LockDaemonProcessForAccept", + blockingPID, + ) if _, err := blockingTx.Exec(ctx, `SELECT pg_sleep(0.3)`); err != nil { t.Fatalf("wait for daemon runtime lease expiry: %v", err) } @@ -380,7 +386,10 @@ func TestSupersededDaemonRuntimeCannotReportGrantedProcessAtStorageBoundary( SourceEndedAt: fixture.Now.Add(4 * time.Second), }, ); !errors.Is(err, storeerr.ErrDaemonRuntimeUnregistered) { - t.Fatalf("superseded runtime after replacement expiry error = %v, want ErrDaemonRuntimeUnregistered", err) + t.Fatalf( + "superseded runtime after replacement expiry error = %v, want ErrDaemonRuntimeUnregistered", + err, + ) } if _, err := fixture.Store.Execution().RegisterDaemonRuntime( ctx, @@ -418,7 +427,13 @@ func TestStartProcessOmitsUnavailableEnvironmentSecrets(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "unavailable_secret_"+scenario) - user := mustCreateProjectDeveloperUser(t, ctx, fixture.Store, "unavailable-secret@example.com", "Secret Tester") + user := mustCreateProjectDeveloperUser( + t, + ctx, + fixture.Store, + "unavailable-secret@example.com", + "Secret Tester", + ) secretIDs := make(map[string]uuid.UUID) secretRefs := make(map[string]string) var grantID uuid.UUID @@ -441,12 +456,13 @@ func TestStartProcessOmitsUnavailableEnvironmentSecrets(t *testing.T) { t.Fatalf("create %s secret: %v", name, err) } if input.OwnerKind == secretstore.SecretOwnerUser { - grant, err := fixture.Store.Secrets().CreateSecretGrant(ctx, secretstore.CreateSecretGrantInput{ - OrgID: testOrgID, - SecretID: secret.ID, - TargetProjectID: testProjectID, - Actor: userPrincipal(user.ID), - }) + grant, err := fixture.Store.Secrets(). + CreateSecretGrant(ctx, secretstore.CreateSecretGrantInput{ + OrgID: testOrgID, + SecretID: secret.ID, + TargetProjectID: testProjectID, + Actor: userPrincipal(user.ID), + }) if err != nil { t.Fatalf("grant secret to project: %v", err) } @@ -468,12 +484,13 @@ func TestStartProcessOmitsUnavailableEnvironmentSecrets(t *testing.T) { t.Helper() var err error if scenario == "grant_revoked" { - _, err = fixture.Store.Secrets().DeleteSecretGrant(ctx, secretstore.DeleteSecretGrantInput{ - OrgID: testOrgID, - SecretID: secretIDs["unavailable"], - GrantID: grantID, - Actor: userPrincipal(user.ID), - }) + _, err = fixture.Store.Secrets(). + DeleteSecretGrant(ctx, secretstore.DeleteSecretGrantInput{ + OrgID: testOrgID, + SecretID: secretIDs["unavailable"], + GrantID: grantID, + Actor: userPrincipal(user.ID), + }) } else { _, err = fixture.Store.Secrets().DeleteSecret(ctx, secretstore.DeleteSecretInput{ OrgID: testOrgID, @@ -489,9 +506,15 @@ func TestStartProcessOmitsUnavailableEnvironmentSecrets(t *testing.T) { removeAccess() } process, err := startProcessForTest(ctx, fixture.Store, executionstore.ExecuteToolCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ToolCallID: createToolCallForProcessTest(t, ctx, fixture, "unavailable_secret_process", "run_command"), + ProjectID: testProjectID, + AgentID: fixture.AgentID, + ToolCallID: createToolCallForProcessTest( + t, + ctx, + fixture, + "unavailable_secret_process", + "run_command", + ), RuntimeLockID: fixture.Lock.ID, }, executionstore.CreateProcessInput{ ExecutionSpec: processcmd.ForShell("echo ok", "", ""), @@ -503,10 +526,11 @@ func TestStartProcessOmitsUnavailableEnvironmentSecrets(t *testing.T) { if scenario == "deleted_after_creation" { removeAccess() } - offers, err := fixture.Store.Execution().ListDaemonProcessOffers(ctx, executionstore.DaemonWorkInput{ - Authority: fixture.authority(), - Limit: 10, - }) + offers, err := fixture.Store.Execution(). + ListDaemonProcessOffers(ctx, executionstore.DaemonWorkInput{ + Authority: fixture.authority(), + Limit: 10, + }) if err != nil { t.Fatalf("list process offers: %v", err) } @@ -517,7 +541,8 @@ func TestStartProcessOmitsUnavailableEnvironmentSecrets(t *testing.T) { if offer.PreparationError != "" || offer.RetryError != nil { t.Fatalf("process preparation failed: %q, %v", offer.PreparationError, offer.RetryError) } - if len(offer.Env) != 2 || offer.Env["APP_ENV"] != "test" || offer.Env["RETAINED"] != "retained-value" { + if len(offer.Env) != 2 || offer.Env["APP_ENV"] != "test" || + offer.Env["RETAINED"] != "retained-value" { t.Fatalf("process environment = %+v, want literal and retained secret only", offer.Env) } }) @@ -575,7 +600,11 @@ func TestStartProcessSnapshotsExecutionConfig(t *testing.T) { } if process.ExecutionSpec.Shell.Command != "echo ok" || process.ExecutionSpec.Shell.Shell != "default" || process.InitialWaitMS != 750 { - t.Fatalf("stored command intent = %q/%q", process.ExecutionSpec.Shell.Command, process.ExecutionSpec.Shell.Shell) + t.Fatalf( + "stored command intent = %q/%q", + process.ExecutionSpec.Shell.Command, + process.ExecutionSpec.Shell.Shell, + ) } offers, err := fixture.Store.Execution().ListDaemonProcessOffers(ctx, executionstore.DaemonWorkInput{ Authority: fixture.authority(), @@ -584,7 +613,8 @@ func TestStartProcessSnapshotsExecutionConfig(t *testing.T) { if err != nil { t.Fatalf("list process offers: %v", err) } - if len(offers) != 1 || offers[0].Env["APP_ENV"] != "test" || offers[0].Env["API_TOKEN"] != "secret-value" || + if len(offers) != 1 || offers[0].Env["APP_ENV"] != "test" || + offers[0].Env["API_TOKEN"] != "secret-value" || offers[0].Env["API_TOKEN_COPY"] != "secret-value" || offers[0].Process.InitialWaitMS != 750 { t.Fatalf("process offers = %+v, want resolved binding environment", offers) @@ -606,12 +636,16 @@ func TestStartProcessSnapshotsExecutionConfig(t *testing.T) { !strings.Contains(offers[0].RetryError.Error(), "secret key wrapper is required") { t.Fatalf("process offers with unavailable secret key wrapper = %+v", offers) } - queuedProcess, err := fixture.Store.Execution().GetProcess(ctx, testProjectID, fixture.AgentID, process.ID) + queuedProcess, err := fixture.Store.Execution(). + GetProcess(ctx, testProjectID, fixture.AgentID, process.ID) if err != nil { t.Fatalf("get process after transient environment resolution failure: %v", err) } if queuedProcess.State != executionstore.ProcessStateQueued { - t.Fatalf("process state after transient environment resolution failure = %q, want queued", queuedProcess.State) + t.Fatalf( + "process state after transient environment resolution failure = %q, want queued", + queuedProcess.State, + ) } nulSecret, _, err := fixture.Store.Secrets().CreateSecret(ctx, secretstore.CreateSecretInput{ OrgID: testOrgID, @@ -675,7 +709,8 @@ func TestStartProcessSnapshotsExecutionConfig(t *testing.T) { if !found { t.Fatal("expected process accept") } - if accept.Process.ExecutionSpec.Shell.Command != "echo ok" || accept.Process.ExecutionSpec.Shell.Shell != "default" || + if accept.Process.ExecutionSpec.Shell.Command != "echo ok" || + accept.Process.ExecutionSpec.Shell.Shell != "default" || accept.Process.InitialWaitMS != 750 { t.Fatalf("accept process = %+v", accept.Process) } @@ -758,7 +793,12 @@ func TestMachineExecutionDefaultUpdatesApplyOnlyToNewProcesses(t *testing.T) { t.Fatalf("replay first process after machine defaults change: %v", err) } if replayed.ID != first.ID || replayed.Cwd != initialCwd { - t.Fatalf("replayed first process = %+v, want original process %s with cwd %q", replayed, first.ID, initialCwd) + t.Fatalf( + "replayed first process = %+v, want original process %s with cwd %q", + replayed, + first.ID, + initialCwd, + ) } second, err := startProcessForTest(ctx, fixture.Store, executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, @@ -1212,10 +1252,10 @@ func TestRunCommandStartCompletesLinkedToolCallWhenAddressable(t *testing.T) { if !strings.Contains(typedBody, publicProcessID) || !strings.Contains(typedBody, `"state": "running"`) { t.Fatalf("typed started tool result missing process handle/state: %s", typedBody) } - var wakeupBefore []byte + var wakeupBefore time.Time if err := fixture.Store.pool.QueryRow(ctx, ` -SELECT wake.metadata -FROM agent_wakeups wake +SELECT wake.logical_ready_at +FROM agent_execution_state wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.project_id = $1 AND wake.agent_id = $2 `, testProjectID, fixture.AgentID).Scan(&wakeupBefore); err != nil { @@ -1234,29 +1274,38 @@ WHERE agent.project_id = $1 AND wake.agent_id = $2 }); err != nil { t.Fatalf("complete already-addressable process: %v", err) } - var wakeupAfter []byte + var wakeupAfter time.Time if err := fixture.Store.pool.QueryRow(ctx, ` -SELECT wake.metadata -FROM agent_wakeups wake +SELECT wake.logical_ready_at +FROM agent_execution_state wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.project_id = $1 AND wake.agent_id = $2 `, testProjectID, fixture.AgentID).Scan(&wakeupAfter); err != nil { t.Fatalf("load terminal-report wakeup: %v", err) } - if !sameJSON(wakeupBefore, wakeupAfter) { + if !wakeupBefore.Equal(wakeupAfter) { t.Fatalf( - "benign terminal report changed wakeup metadata before=%s after=%s", - string(wakeupBefore), - string(wakeupAfter), + "benign terminal report changed readiness before=%s after=%s", + wakeupBefore, + wakeupAfter, ) } + if count := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); count != 0 { + t.Fatalf("owned agent has %d wakeups", count) + } } func TestFileTransferWaitsForTerminalResult(t *testing.T) { t.Parallel() - memoryOutput := `{"path":"/memory/team/screenshot.png","digest":"sha256:` + strings.Repeat("a", 64) + `"}` + "\n" + memoryOutput := `{"path":"/memory/team/screenshot.png","digest":"sha256:` + strings.Repeat( + "a", + 64, + ) + `"}` + "\n" downloadOutput := `{"digest":"sha256:` + strings.Repeat("a", 64) + `"}` - conflict := `{"code":"file_content_conflict","error":"file changed","current_digest":"sha256:` + strings.Repeat("b", 64) + `"}` + conflict := `{"code":"file_content_conflict","error":"file changed","current_digest":"sha256:` + strings.Repeat( + "b", + 64, + ) + `"}` const memoryFailureOutput = "upload file: memory changed; read it and retry\n" for _, test := range []struct { name string @@ -1295,25 +1344,37 @@ func TestFileTransferWaitsForTerminalResult(t *testing.T) { { name: "artifact_missing_metadata", toolName: "upload_file", direction: processcmd.FileTransferUpload, createArtifact: true, - input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), - wantError: "upload completed without valid file metadata", wantOutcome: executionstore.ToolResultOutcomeFailed, + input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), + wantError: "upload completed without valid file metadata", + wantOutcome: executionstore.ToolResultOutcomeFailed, }, { - name: "artifact_wrong_path", toolName: "upload_file", direction: processcmd.FileTransferUpload, createArtifact: true, - input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), artifactReport: "wrong_path", - wantError: "upload completed without valid file metadata", wantOutcome: executionstore.ToolResultOutcomeFailed, + name: "artifact_wrong_path", + toolName: "upload_file", + direction: processcmd.FileTransferUpload, + createArtifact: true, + input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), + artifactReport: "wrong_path", + wantError: "upload completed without valid file metadata", + wantOutcome: executionstore.ToolResultOutcomeFailed, }, { name: "artifact_wrong_digest", toolName: "upload_file", direction: processcmd.FileTransferUpload, createArtifact: true, input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), artifactReport: "wrong_digest", - wantError: "upload completed without valid file metadata", wantOutcome: executionstore.ToolResultOutcomeFailed, + wantError: "upload completed without valid file metadata", + wantOutcome: executionstore.ToolResultOutcomeFailed, }, { - name: "unrelated_artifact", toolName: "upload_file", direction: processcmd.FileTransferUpload, createArtifact: true, - input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), artifactReport: "valid", + name: "unrelated_artifact", + toolName: "upload_file", + direction: processcmd.FileTransferUpload, + createArtifact: true, + input: json.RawMessage(`{"path":"/artifacts","source":"screenshot.png"}`), + artifactReport: "valid", unrelatedArtifact: true, - wantError: "upload completed without an artifact", wantOutcome: executionstore.ToolResultOutcomeFailed, + wantError: "upload completed without an artifact", + wantOutcome: executionstore.ToolResultOutcomeFailed, }, { name: "memory", @@ -1364,7 +1425,9 @@ func TestFileTransferWaitsForTerminalResult(t *testing.T) { name: "fast_artifact_download", toolName: "download_file", direction: processcmd.FileTransferDownload, skipStart: true, input: mustTestRawJSON(t, map[string]string{ - "path": "/artifacts/" + publicResourceID(publicid.KindArtifact, uuid.New()), "destination": "screenshot.png", + "path": "/artifacts/" + publicResourceID(publicid.KindArtifact, + uuid.New()), + "destination": "screenshot.png", }), terminalResult: mustTestRawJSON(t, map[string]any{ "output": "", "cursor": 0, "next_cursor": 0, @@ -1415,7 +1478,10 @@ func TestFileTransferWaitsForTerminalResult(t *testing.T) { storeID := uuid.New() storeName := "transfer-" + strings.ReplaceAll(storeID.String(), "-", "") if _, err := fixture.Store.pool.Exec(ctx, - `INSERT INTO memory_stores(id, project_id, name) VALUES ($1, $2, $3)`, storeID, testProjectID, storeName); err != nil { + `INSERT INTO memory_stores(id, project_id, name) VALUES ($1, $2, $3)`, + storeID, + testProjectID, + storeName); err != nil { t.Fatal(err) } test.terminalResult = json.RawMessage(strings.ReplaceAll( @@ -1487,7 +1553,9 @@ func TestFileTransferWaitsForTerminalResult(t *testing.T) { ID: process.ID, Authority: fixture.authority(), SourceStartedAt: fixture.Now.Add(time.Second), - Result: json.RawMessage(`{"output":"warning\n","cursor":0,"next_cursor":8,"truncated":false}`), + Result: json.RawMessage( + `{"output":"warning\n","cursor":0,"next_cursor":8,"truncated":false}`, + ), }, ) if err != nil { @@ -1534,13 +1602,15 @@ func TestFileTransferWaitsForTerminalResult(t *testing.T) { } completionInput := executionstore.CompleteDaemonProcessInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ID: process.ID, - Authority: fixture.authority(), - State: executionstore.ProcessStateExited, - ExitCode: &test.exitCode, - Result: json.RawMessage(`{"output":"not json","cursor":0,"next_cursor":100,"truncated":true}`), + ProjectID: testProjectID, + AgentID: fixture.AgentID, + ID: process.ID, + Authority: fixture.authority(), + State: executionstore.ProcessStateExited, + ExitCode: &test.exitCode, + Result: json.RawMessage( + `{"output":"not json","cursor":0,"next_cursor":100,"truncated":true}`, + ), SourceStartedAt: fixture.Now.Add(time.Second), SourceEndedAt: fixture.Now.Add(2 * time.Second), } @@ -1691,7 +1761,11 @@ WHERE result.agent_id = $1 AND result.tool_call_id = $2 } } if !foundCompactionResult { - t.Fatalf("upload artifact compaction result = %+v, want %s", compactionEvents, wantModelContent) + t.Fatalf( + "upload artifact compaction result = %+v, want %s", + compactionEvents, + wantModelContent, + ) } } else { var wantValue map[string]any @@ -1807,7 +1881,10 @@ func TestRunCommandTerminalResultRetainsCanonicalProcessHandle(t *testing.T) { publicProcessID := publicResourceID(publicid.KindProcess, process.ID) if !strings.Contains(string(completed.ResultContentParts), publicProcessID) || strings.Contains(string(completed.ResultContentParts), "prc_untrusted") { - t.Fatalf("terminal run_command result has the wrong process handle: %s", string(completed.ResultContentParts)) + t.Fatalf( + "terminal run_command result has the wrong process handle: %s", + string(completed.ResultContentParts), + ) } } @@ -1948,7 +2025,11 @@ func TestProcessAcceptUsesReplacementGrantAfterGrantRotation(t *testing.T) { t.Fatalf("start process: %v", err) } if process.AgentMachineBindingID != fixture.BindingID { - t.Fatalf("process binding = %s, want fixture binding %s", process.AgentMachineBindingID, fixture.BindingID) + t.Fatalf( + "process binding = %s, want fixture binding %s", + process.AgentMachineBindingID, + fixture.BindingID, + ) } if _, err := fixture.Store.Execution().DeleteProjectMachineGrant( ctx, @@ -1961,15 +2042,23 @@ func TestProcessAcceptUsesReplacementGrantAfterGrantRotation(t *testing.T) { if err != nil { t.Fatalf("get revoked-grant process: %v", err) } - if current.State != executionstore.ProcessStateFailed || current.StateReasonCode != "project_machine_grant_revoked" || + if current.State != executionstore.ProcessStateFailed || + current.StateReasonCode != "project_machine_grant_revoked" || current.SourceEndedAt != nil { t.Fatalf("revoked-grant process = %+v, want failed/project_machine_grant_revoked", current) } - toolCall, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, originalToolCallID) + toolCall, err := fixture.Store.Execution(). + GetToolCall(ctx, testProjectID, fixture.AgentID, originalToolCallID) if err != nil { t.Fatalf("get revoked-grant tool call: %v", err) } - assertCompletedToolCallWithResult(t, fixture.Store, fixture.AgentID, toolCall, "project_machine_grant_revoked") + assertCompletedToolCallWithResult( + t, + fixture.Store, + fixture.AgentID, + toolCall, + "project_machine_grant_revoked", + ) newGrant, _, err := fixture.Store.Execution().CreateProjectMachineGrant( ctx, executionstore.CreateProjectMachineGrantInput{ @@ -1998,7 +2087,11 @@ func TestProcessAcceptUsesReplacementGrantAfterGrantRotation(t *testing.T) { t.Fatalf("start replacement-grant process: %v", err) } if replacementProcess.AgentMachineBindingID != fixture.BindingID { - t.Fatalf("replacement process binding = %s, want %s", replacementProcess.AgentMachineBindingID, fixture.BindingID) + t.Fatalf( + "replacement process binding = %s, want %s", + replacementProcess.AgentMachineBindingID, + fixture.BindingID, + ) } accept, found, err := acceptDaemonProcessForTest( ctx, @@ -2067,7 +2160,8 @@ func TestRevokeProjectMachineGrantTerminatesActiveDaemonProcess(t *testing.T) { if err != nil { t.Fatalf("get revoked-grant process: %v", err) } - if current.State != executionstore.ProcessStateUnknown || current.StateReasonCode != "project_machine_grant_revoked" || + if current.State != executionstore.ProcessStateUnknown || + current.StateReasonCode != "project_machine_grant_revoked" || current.SourceEndedAt != nil { t.Fatalf("revoked-grant process = %+v, want unknown/project_machine_grant_revoked", current) } @@ -2315,6 +2409,13 @@ WHERE call.project_id = $1 t.Fatalf("load tool call state: %v", err) } if gotType != wantType || gotState != wantState { - t.Fatalf("tool call %s type/state = %s/%s, want %s/%s", toolCallID, gotType, gotState, wantType, wantState) + t.Fatalf( + "tool call %s type/state = %s/%s, want %s/%s", + toolCallID, + gotType, + gotState, + wantType, + wantState, + ) } } diff --git a/internal/storage/executionstore/process_runtime_lease_recovery_integration_test.go b/internal/storage/executionstore/process_runtime_lease_recovery_integration_test.go index a003c10e6..8772f32ce 100644 --- a/internal/storage/executionstore/process_runtime_lease_recovery_integration_test.go +++ b/internal/storage/executionstore/process_runtime_lease_recovery_integration_test.go @@ -18,7 +18,6 @@ import ( "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/storage" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/testutil/integrationdb" ) @@ -47,13 +46,17 @@ func TestRuntimeLockReaperCrossesBatchBoundary(t *testing.T) { if err != nil { t.Fatalf("create maintenance cursor agent %d: %v", index, err) } - if _, _, _, err := store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: agent.ID, - Actor: mustOmnaraActorParams(t, user.ID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"maintenance cursor work"}]`), - IdempotencyKey: fmt.Sprintf("maintenance-cursor-input-%03d", index), - }); err != nil { + if _, + _, + _, + err := store.Execution().CreateAgentContentInput(ctx, + executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: agent.ID, + Actor: mustOmnaraActorParams(t, user.ID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"maintenance cursor work"}]`), + IdempotencyKey: fmt.Sprintf("maintenance-cursor-input-%03d", index), + }); err != nil { t.Fatalf("create maintenance cursor input %d: %v", index, err) } } @@ -113,14 +116,14 @@ func assertRuntimeLockReapCounts( t.Helper() var wakeups, locks int if err := pool.QueryRow(ctx, ` -SELECT count(*) FILTER (WHERE wake.metadata->>'reason' = $2), +SELECT count(*), (SELECT count(*) FROM agent_runtime_locks runtime_lock JOIN agents lock_agent ON lock_agent.id = runtime_lock.agent_id WHERE lock_agent.project_id = $1) FROM agent_wakeups wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.project_id = $1 -`, projectID, "runtime_lock_reap").Scan(&wakeups, &locks); err != nil { +`, projectID).Scan(&wakeups, &locks); err != nil { t.Fatalf("count runtime-lock reap state: %v", err) } if wakeups != wantWakeups || locks != wantLocks { @@ -146,13 +149,14 @@ func TestClaimNextAgentWorkRecoversUnstartedTurnAfterRuntimeRelease(t *testing.T ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"recover me"}]`), - IdempotencyKey: "claim-recovers-unstarted-turn", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"recover me"}]`), + IdempotencyKey: "claim-recovers-unstarted-turn", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -162,7 +166,12 @@ func TestClaimNextAgentWorkRecoversUnstartedTurnAfterRuntimeRelease(t *testing.T } if !found || firstClaim.Kind != executionstore.AgentWorkModel || firstClaim.Model.TurnID == uuid.Nil || !claimedOpeningInputIDsEqual(firstClaim, input.ID) { - t.Fatalf("initial claim = %+v found=%v, want executable turn for input %s", firstClaim, found, input.ID) + t.Fatalf( + "initial claim = %+v found=%v, want executable turn for input %s", + firstClaim, + found, + input.ID, + ) } if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( ctx, @@ -186,7 +195,8 @@ func TestClaimNextAgentWorkRecoversUnstartedTurnAfterRuntimeRelease(t *testing.T if !found || recoveredClaim.Kind != executionstore.AgentWorkModel { t.Fatalf("recovered claim = %+v found=%v, want executable", recoveredClaim, found) } - if recoveredClaim.Model.TurnID != firstClaim.Model.TurnID || !claimedOpeningInputIDsEqual(recoveredClaim, input.ID) { + if recoveredClaim.Model.TurnID != firstClaim.Model.TurnID || + !claimedOpeningInputIDsEqual(recoveredClaim, input.ID) { t.Fatalf("recovered claim = %+v, want same turn/input", recoveredClaim) } } @@ -204,13 +214,14 @@ func TestStartedModelContextOnLiveLeaseIsNotSchedulerResumable(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"active context"}]`), - IdempotencyKey: "claim-keeps-active-context-claimable", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"active context"}]`), + IdempotencyKey: "claim-keeps-active-context-claimable", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -218,7 +229,8 @@ func TestStartedModelContextOnLiveLeaseIsNotSchedulerResumable(t *testing.T) { if err != nil { t.Fatalf("claim input work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) { t.Fatalf("claim = %+v found=%v, want executable input %s", claim, found, input.ID) } claimTestNormalModelCallForWork(t, ctx, fixture, claim, now.Add(2*time.Second)) @@ -242,7 +254,7 @@ func TestClaimNextAgentWorkRecoversFailedAbandonedContext(t *testing.T) { fixture := newProcessDaemonFixture(t, ctx, "claim_recovers_failed_abandoned_context") fixture.Store = newIntegrationStore( fixture.Store.pool, - storage.WithModelCallRetryBackoff(executionstore.ModelCallRetryBackoff), + storage.WithModelCallRetryBackoff(modelprotocol.RetryBackoff), ) now := fixture.Now.Add(time.Minute) if err := fixture.Store.Execution().ReleaseAgentRuntimeLock( @@ -253,13 +265,14 @@ func TestClaimNextAgentWorkRecoversFailedAbandonedContext(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"retry abandoned"}]`), - IdempotencyKey: "claim-recovers-failed-abandoned-context", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"retry abandoned"}]`), + IdempotencyKey: "claim-recovers-failed-abandoned-context", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -267,7 +280,8 @@ func TestClaimNextAgentWorkRecoversFailedAbandonedContext(t *testing.T) { if err != nil { t.Fatalf("claim input work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) { t.Fatalf("claim = %+v found=%v, want executable input %s", claim, found, input.ID) } modelClaim := claimTestNormalModelCallForWork(t, ctx, fixture, claim, now.Add(2*time.Second)) @@ -304,12 +318,16 @@ func TestClaimNextAgentWorkRecoversFailedAbandonedContext(t *testing.T) { contextRecord, ) } - wantRetryAt := contextRecord.CompletedAt.Add(executionstore.ModelCallRetryBackoff( + wantRetryAt := contextRecord.CompletedAt.Add(modelprotocol.RetryBackoff( contextRecord.AttemptNumber, contextRecord.ID.String(), )) if !contextRecord.RetryAt.Equal(wantRetryAt) { - t.Fatalf("interrupted retry_at = %s, want exact durable backoff %s", contextRecord.RetryAt, wantRetryAt) + t.Fatalf( + "interrupted retry_at = %s, want exact durable backoff %s", + contextRecord.RetryAt, + wantRetryAt, + ) } var wakeupAt time.Time if err := fixture.Store.pool.QueryRow(ctx, ` @@ -324,10 +342,11 @@ WHERE agent.project_id = $1 AND wake.agent_id = $2 t.Fatalf("interrupted wakeup = %s, want retry_at %s", wakeupAt, contextRecord.RetryAt) } requireAgentWakeupCoverage(t, ctx, fixture.Store, testProjectID, fixture.AgentID) - if turn, err := fixture.Store.q.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: testProjectID, AgentID: fixture.AgentID}, - ); err != nil { + if turn, + err := currentContinuableTurnForTest(ctx, + fixture.Store, + testProjectID, + fixture.AgentID); err != nil { t.Fatalf("load current continuable turn for recoverable failed context: %v", err) } else if turn.ID != claim.Model.TurnID { t.Fatalf("recoverable failed context current turn = %+v, want turn %s", turn, claim.Model.TurnID) @@ -371,9 +390,14 @@ WHERE agent.id = wake.agent_id AND agent.project_id = $1 AND wake.agent_id = $2` if err != nil { t.Fatalf("claim failed abandoned context recovery: %v", err) } - if !found || recovered.Kind != executionstore.AgentWorkModel || recovered.Model.TurnID != claim.Model.TurnID || + if !found || recovered.Kind != executionstore.AgentWorkModel || + recovered.Model.TurnID != claim.Model.TurnID || !claimedOpeningInputIDsEqual(recovered, input.ID) { - t.Fatalf("failed abandoned context recovery claim = %+v found=%v, want same executable turn/input", recovered, found) + t.Fatalf( + "failed abandoned context recovery claim = %+v found=%v, want same executable turn/input", + recovered, + found, + ) } retryClaim, err := fixture.Store.Execution().ClaimNextModelCallContext( ctx, @@ -425,13 +449,14 @@ func TestClaimNextAgentWorkRecoversAmbiguousAbandonedContext(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"retry ambiguous send"}]`), - IdempotencyKey: "claim-recovers-ambiguous-abandoned-context", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"retry ambiguous send"}]`), + IdempotencyKey: "claim-recovers-ambiguous-abandoned-context", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -439,7 +464,8 @@ func TestClaimNextAgentWorkRecoversAmbiguousAbandonedContext(t *testing.T) { if err != nil { t.Fatalf("claim input work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) { t.Fatalf("claim = %+v found=%v, want executable input %s", claim, found, input.ID) } modelClaim := claimTestNormalModelCallForWork(t, ctx, fixture, claim, now.Add(2*time.Second)) @@ -484,7 +510,8 @@ func TestClaimNextAgentWorkRecoversAmbiguousAbandonedContext(t *testing.T) { if err != nil { t.Fatalf("claim ambiguous abandoned context recovery: %v", err) } - if !found || recovered.Kind != executionstore.AgentWorkModel || recovered.Model.TurnID != claim.Model.TurnID || + if !found || recovered.Kind != executionstore.AgentWorkModel || + recovered.Model.TurnID != claim.Model.TurnID || !claimedOpeningInputIDsEqual(recovered, input.ID) { t.Fatalf("ambiguous recovery claim = %+v found=%v, want same executable turn/input", recovered, found) } @@ -617,7 +644,13 @@ func TestReapedModelCallWorkerCannotPublishAfterReplacementClaim(t *testing.T) { ); err != nil || found { t.Fatalf("stale worker output found=%v err=%v, want none", found, err) } - assertModelCallContextState(t, ctx, fixture, oldClaim.Context.ID, executionstore.ModelCallContextFailed) + assertModelCallContextState( + t, + ctx, + fixture, + oldClaim.Context.ID, + executionstore.ModelCallContextFailed, + ) assertReplacementModelCallContextUnchanged(t, ctx, fixture, replacementClaim) }) @@ -626,16 +659,17 @@ func TestReapedModelCallWorkerCannotPublishAfterReplacementClaim(t *testing.T) { ctx := context.Background() fixture, oldWork, oldClaim, replacementWork, replacementClaim := claimReplacementAfterReapedCompaction(t, ctx, "stale_worker_checkpoint") - _, err := fixture.Store.Execution().PublishContextCheckpoint(ctx, executionstore.PublishContextCheckpointInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: oldWork.RuntimeLock.ID, - ModelCallContextID: oldClaim.Context.ID, - Summary: "late checkpoint", - APIFormat: modelprotocol.APIFormatOpenAIResponses, - APIVariant: modelprotocol.APIVariantDefault, - ProviderResponseID: "resp_stale_checkpoint", - }) + _, err := fixture.Store.Execution(). + PublishContextCheckpoint(ctx, executionstore.PublishContextCheckpointInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: oldWork.RuntimeLock.ID, + ModelCallContextID: oldClaim.Context.ID, + Summary: "late checkpoint", + APIFormat: modelprotocol.APIFormatOpenAIResponses, + APIVariant: modelprotocol.APIVariantDefault, + ProviderResponseID: "resp_stale_checkpoint", + }) if !errors.Is(err, storeerr.ErrRuntimeLockInactive) { t.Fatalf("stale checkpoint publication error = %v, want %v", err, storeerr.ErrRuntimeLockInactive) } @@ -647,7 +681,13 @@ func TestReapedModelCallWorkerCannotPublishAfterReplacementClaim(t *testing.T) { ); err != nil || found { t.Fatalf("stale worker checkpoint found=%v err=%v, want none", found, err) } - assertModelCallContextState(t, ctx, fixture, oldClaim.Context.ID, executionstore.ModelCallContextFailed) + assertModelCallContextState( + t, + ctx, + fixture, + oldClaim.Context.ID, + executionstore.ModelCallContextFailed, + ) assertReplacementModelCallContextUnchanged(t, ctx, fixture, replacementClaim) checkpoint, err := fixture.Store.Execution().PublishContextCheckpoint( ctx, @@ -672,8 +712,20 @@ func TestReapedModelCallWorkerCannotPublishAfterReplacementClaim(t *testing.T) { replacementClaim.Context.ID, ) } - assertModelCallContextState(t, ctx, fixture, oldClaim.Context.ID, executionstore.ModelCallContextFailed) - assertModelCallContextState(t, ctx, fixture, replacementClaim.Context.ID, executionstore.ModelCallContextSucceeded) + assertModelCallContextState( + t, + ctx, + fixture, + oldClaim.Context.ID, + executionstore.ModelCallContextFailed, + ) + assertModelCallContextState( + t, + ctx, + fixture, + replacementClaim.Context.ID, + executionstore.ModelCallContextSucceeded, + ) contextRecord, found, err := fixture.Store.Execution().GetModelCallContext( ctx, testProjectID, @@ -726,18 +778,25 @@ func TestCompleteRuntimeToolCallRejectsCanceledRuntime(t *testing.T) { ); err != nil { t.Fatalf("cancel runtime: %v", err) } - contentParts, err := executionstore.ToolResultContentParts(json.RawMessage(`{"ok":false,"error":"runtime canceled"}`)) + contentParts, err := executionstore.ToolResultContentParts( + json.RawMessage(`{"ok":false,"error":"runtime canceled"}`), + ) if err != nil { t.Fatalf("build runtime failure content parts: %v", err) } - if _, err := fixture.Store.Execution().CompleteRuntimeToolCall(ctx, executionstore.CompleteRuntimeToolCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ID: toolCallID, - RuntimeLockID: fixture.Lock.ID, - Outcome: executionstore.ToolResultOutcomeFailed, - ResultContentParts: contentParts, - }); !errors.Is(err, storeerr.ErrRuntimeLockInactive) { + if _, + err := fixture.Store.Execution().CompleteRuntimeToolCall(ctx, + executionstore.CompleteRuntimeToolCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + ID: toolCallID, + RuntimeLockID: fixture.Lock.ID, + Outcome: executionstore.ToolResultOutcomeFailed, + ResultContentParts: contentParts, + }); !errors.Is( + err, + storeerr.ErrRuntimeLockInactive, + ) { t.Fatalf("complete runtime tool call after cancel error = %v, want runtime lock inactive", err) } var wakeups int @@ -747,8 +806,8 @@ FROM agent_wakeups wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.project_id = $1 AND wake.agent_id = $2 - AND wake.metadata->>'tool_call_id' = $3 -`, testProjectID, fixture.AgentID, toolCallID.String()).Scan(&wakeups); err != nil { + +`, testProjectID, fixture.AgentID).Scan(&wakeups); err != nil { t.Fatalf("count runtime completion wakeups: %v", err) } if wakeups != 0 { @@ -878,14 +937,19 @@ func TestExpiredRuntimeFailsOwnedToolAndFencesLateCompletion(t *testing.T) { if err != nil { t.Fatalf("build late runtime tool result: %v", err) } - if _, err := fixture.Store.Execution().CompleteRuntimeToolCall(ctx, executionstore.CompleteRuntimeToolCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ID: toolCallID, - RuntimeLockID: fixture.Lock.ID, - Outcome: executionstore.ToolResultOutcomeSucceeded, - ResultContentParts: lateParts, - }); !errors.Is(err, storeerr.ErrRuntimeLockInactive) { + if _, + err := fixture.Store.Execution().CompleteRuntimeToolCall(ctx, + executionstore.CompleteRuntimeToolCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + ID: toolCallID, + RuntimeLockID: fixture.Lock.ID, + Outcome: executionstore.ToolResultOutcomeSucceeded, + ResultContentParts: lateParts, + }); !errors.Is( + err, + storeerr.ErrRuntimeLockInactive, + ) { t.Fatalf("late runtime tool completion error = %v, want runtime lock inactive", err) } replacement, err := fixture.Store.Execution().AcquireAgentRuntimeLock( @@ -969,11 +1033,13 @@ func TestExpiredRuntimePreservesDurablyWaitingQuestionInteraction(t *testing.T) found, ) } - if _, err := fixture.Store.Execution().ResolveAgentInteraction(ctx, executionstore.ResolveAgentInteractionInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, ID: interaction.ID, - Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{OptionIndices: []int{0}}}}, - Actor: mustOmnaraActorParams(t, fixture.UserID), - }); err != nil { + if _, + err := fixture.Store.Execution().ResolveAgentInteraction(ctx, + executionstore.ResolveAgentInteractionInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, ID: interaction.ID, + Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{OptionIndices: []int{0}}}}, + Actor: mustOmnaraActorParams(t, fixture.UserID), + }); err != nil { t.Fatalf("answer question after worker interruption: %v", err) } completed, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, toolCallID) diff --git a/internal/storage/executionstore/process_start_cancel_integration_helpers_test.go b/internal/storage/executionstore/process_start_cancel_integration_helpers_test.go index 1725e5c64..6c251d415 100644 --- a/internal/storage/executionstore/process_start_cancel_integration_helpers_test.go +++ b/internal/storage/executionstore/process_start_cancel_integration_helpers_test.go @@ -13,7 +13,6 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) func claimReplacementAfterReapedNormalModelCall( @@ -50,13 +49,14 @@ func claimReplacementAfterReapedNormalModelCallSetup( t.Fatalf("release initial runtime lock: %v", err) } now := fixture.Now.Add(time.Minute) - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise stale worker fencing"}]`), - IdempotencyKey: fixtureName, - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise stale worker fencing"}]`), + IdempotencyKey: fixtureName, + }) if err != nil { t.Fatalf("create stale-worker input: %v", err) } @@ -152,12 +152,13 @@ func claimReplacementModelCallContext( work.Model.ModelCallContextID != predecessorContextID { t.Fatalf("claim replacement model work = %+v found=%v err=%v", work, found, err) } - claim, err := fixture.Store.Execution().ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - PredecessorModelCallContextID: predecessorContextID, - RuntimeLockID: work.RuntimeLock.ID, - }) + claim, err := fixture.Store.Execution(). + ClaimNextModelCallContext(ctx, executionstore.ClaimNextModelCallContextInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + PredecessorModelCallContextID: predecessorContextID, + RuntimeLockID: work.RuntimeLock.ID, + }) if err != nil { t.Fatalf("claim replacement model context: %v", err) } @@ -242,13 +243,14 @@ func claimExhaustedCompactionModelContext( ); err != nil { t.Fatalf("release compaction fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise compaction retry limit"}]`), - IdempotencyKey: fixtureName, - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise compaction retry limit"}]`), + IdempotencyKey: fixtureName, + }) if err != nil { t.Fatalf("create compaction content input: %v", err) } @@ -397,13 +399,14 @@ func claimExhaustedNormalModelContext( ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise model call retry limit"}]`), - IdempotencyKey: fixtureName, - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"exercise model call retry limit"}]`), + IdempotencyKey: fixtureName, + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -655,7 +658,7 @@ func claimTestNormalModelCallForWork( now time.Time, ) executionstore.ModelCallClaim { t.Helper() - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent for model context: %v", err) } @@ -668,14 +671,15 @@ func claimTestNormalModelCallForWork( Scan(&openingEventSequence); err != nil { t.Fatalf("load opening event sequence: %v", err) } - modelClaim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: claim.RuntimeLock.ID, - OpeningInputIDs: claim.Model.InputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: openingEventSequence, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: claim.RuntimeLock.ID, + OpeningInputIDs: claim.Model.InputIDs, + }) + modelClaim := prepared1.Claim + if err != nil { t.Fatalf("claim model context: %v", err) } @@ -705,7 +709,8 @@ func requireAgentWakeupCoverage( t.Helper() var runnable, hasWakeup, hasRuntimeLock bool if err := store.pool.QueryRow(ctx, ` -SELECT agent_next_wakeup_ready_at($1, $2) IS NOT NULL, +SELECT EXISTS(SELECT 1 FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id +WHERE a.project_id=$1 AND a.id=$2 AND h.logical_ready_at IS NOT NULL), EXISTS (SELECT 1 FROM agent_wakeups WHERE agent_id = $2), EXISTS (SELECT 1 FROM agent_runtime_locks WHERE agent_id = $2) `, projectID, agentID).Scan(&runnable, &hasWakeup, &hasRuntimeLock); err != nil { @@ -807,18 +812,19 @@ func retryBackoffWithQueuedInput( if !found || len(admitted.Events) != 1 { t.Fatalf("admit retry opening input found=%v turn=%+v", found, admitted) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load retry agent: %v", err) } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{openingInput.ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) + prepared2, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{openingInput.ID}, + }) + claim := prepared2.Claim + if err != nil { t.Fatalf("claim retry model call: %v", err) } @@ -856,6 +862,12 @@ func retryBackoffWithQueuedInput( if err != nil { t.Fatalf("create queued input during retry backoff: %v", err) } + if err := fixture.Store.Execution().ReleaseAgentRuntimeLock(ctx, + testProjectID, + fixture.AgentID, + fixture.Lock.ID); err != nil { + t.Fatal(err) + } if readyAt := agentWakeupReadyAt(t, ctx, fixture.Store, fixture.AgentID); !readyAt.Equal(retryAt) { t.Fatalf("initial retry wakeup = %s, want %s", readyAt, retryAt) } @@ -889,15 +901,38 @@ func requestAgentRuntimeCancelForTest( ctx context.Context, store *Store, projectID, agentID uuid.UUID, - now time.Time, + _ time.Time, ) (executionstore.AgentRuntimeLockRecord, error) { - row, err := dbsqlc.New(store.pool). - RequestAgentRuntimeCancel(ctx, dbsqlc.RequestAgentRuntimeCancelParams{ProjectID: projectID, AgentID: agentID}) + var r executionstore.AgentRuntimeLockRecord + err := store.pool.QueryRow(ctx, + `UPDATE agent_runtime_locks r SET cancel_requested_at=coalesce(cancel_requested_at,statement_timestamp()) FROM agents a WHERE a.id=r.agent_id AND a.project_id=$1 AND a.id=$2 AND r.cancel_requested_at IS NULL RETURNING r.id,r.agent_id,r.worker_process_id,r.started_at,r.renewed_at,r.lease_expires_at,r.cancel_requested_at`, + projectID, + agentID). + Scan(&r.ID, + &r.AgentID, + &r.WorkerProcessID, + &r.StartedAt, + &r.RenewedAt, + &r.LeaseExpiresAt, + &r.CancelRequestedAt) if errors.Is(err, pgx.ErrNoRows) { return executionstore.AgentRuntimeLockRecord{}, nil } - if err != nil { - return executionstore.AgentRuntimeLockRecord{}, err - } - return executionstore.IntegrationAgentRuntimeLockRecordFromCancelSQLC(row), nil + return r, err +} + +func currentContinuableTurnForTest( + ctx context.Context, + store *Store, + projectID, agentID uuid.UUID, +) (executionstore.AgentTurnRecord, error) { + var r executionstore.AgentTurnRecord + err := store.pool.QueryRow(ctx, + `SELECT t.id,t.turn_sequence,t.latest_event_id,t.latest_semantic_event_id FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id JOIN agent_turns t ON t.agent_id=h.agent_id AND t.id=h.current_turn_id WHERE a.project_id=$1 AND a.id=$2 AND h.turn_continuable`, + projectID, + agentID). + Scan(&r.ID, &r.TurnSequence, &r.LatestEventID, &r.LatestSemanticEventID) + r.ProjectID = projectID + r.AgentID = agentID + return r, err } diff --git a/internal/storage/executionstore/process_tool_unreachable_store.go b/internal/storage/executionstore/process_tool_unreachable_store.go index 1ce51bd98..dd7d76308 100644 --- a/internal/storage/executionstore/process_tool_unreachable_store.go +++ b/internal/storage/executionstore/process_tool_unreachable_store.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -88,7 +89,10 @@ func (s *Store) expireProcessToolCallsForMachine( }, ) if err != nil { - return total, fmt.Errorf("list machine-unreachable accepted process tool calls for machine: %w", err) + return total, fmt.Errorf( + "list machine-unreachable accepted process tool calls for machine: %w", + err, + ) } queuedActions, err := s.q.ListMachineUnreachableQueuedProcessActionToolCallsForMachine( ctx, @@ -100,7 +104,10 @@ func (s *Store) expireProcessToolCallsForMachine( }, ) if err != nil { - return total, fmt.Errorf("list machine-unreachable queued process action tool calls for machine: %w", err) + return total, fmt.Errorf( + "list machine-unreachable queued process action tool calls for machine: %w", + err, + ) } acceptedActions, err := s.q.ListMachineUnreachableAcceptedProcessActionToolCallsForMachine( ctx, @@ -112,7 +119,10 @@ func (s *Store) expireProcessToolCallsForMachine( }, ) if err != nil { - return total, fmt.Errorf("list machine-unreachable accepted process action tool calls for machine: %w", err) + return total, fmt.Errorf( + "list machine-unreachable accepted process action tool calls for machine: %w", + err, + ) } if len(queuedProcesses) == 0 && len(acceptedProcesses) == 0 && len(queuedActions) == 0 && len(acceptedActions) == 0 { @@ -182,12 +192,12 @@ func (s *Store) expireQueuedProcessToolCall( process ProcessRecord, graceSeconds int32, ) (bool, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return false, fmt.Errorf("begin queued process expiry: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() qtx := dbsqlc.New(tx) unreachable, err := machineStillUnreachableForToolExpiryTx( ctx, @@ -200,13 +210,10 @@ func (s *Store) expireQueuedProcessToolCall( if err != nil { return false, err } - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ - ProjectID: process.ProjectID, - ID: process.AgentID, - }, - ); err != nil { + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: process.ProjectID, + ID: process.AgentID, + }, agentexecution.ExternalAuthority{}); err != nil { return false, fmt.Errorf("lock agent for queued process expiry: %w", err) } row, err := qtx.ExpireQueuedProcessToolCall( @@ -222,12 +229,7 @@ func (s *Store) expireQueuedProcessToolCall( }, ) if errors.Is(err, pgx.ErrNoRows) { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "missed queued process expiry", - ); err != nil { + if err := unit.Commit(ctx, "missed queued process expiry"); err != nil { return false, err } return false, nil @@ -236,23 +238,19 @@ func (s *Store) expireQueuedProcessToolCall( return false, fmt.Errorf("mark queued process failed: %w", err) } record := processRecordFromSQLC(row) - if err := completeProcessToolCallFromRecordTx( - ctx, - txNotifications, - tx, - qtx, - record, - nil, - record.StateReasonCode, - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return false, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "queued process expiry", - ); err != nil { + if record.ToolCallID != uuid.Nil { + if _, + err := h.CompleteProcess(ctx, + agentexecution.ProcessResult{ID: record.ID, + Observed: nil}); err != nil { + return false, err + } + } + if err := unit.Commit(ctx, "queued process expiry"); err != nil { return false, err } return true, nil @@ -294,12 +292,12 @@ func (s *Store) failMachineUnreachableQueuedProcessActions( action ProcessActionRecord, graceSeconds int32, ) (int64, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return 0, fmt.Errorf("begin machine-unreachable queued action expiry: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() qtx := dbsqlc.New(tx) unreachable, err := machineStillUnreachableForToolExpiryTx( ctx, @@ -313,23 +311,15 @@ func (s *Store) failMachineUnreachableQueuedProcessActions( return 0, err } if !unreachable { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "skipped machine-unreachable queued action expiry", - ); err != nil { + if err := unit.Commit(ctx, "skipped machine-unreachable queued action expiry"); err != nil { return 0, err } return 0, nil } - if _, err := qtx.LockAgentInProject( - ctx, - dbsqlc.LockAgentInProjectParams{ - ProjectID: action.ProjectID, - ID: action.AgentID, - }, - ); err != nil { + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: action.ProjectID, + ID: action.AgentID, + }, agentexecution.ExternalAuthority{}); err != nil { return 0, fmt.Errorf("lock agent for machine-unreachable queued action expiry: %w", err) } rows, err := qtx.MarkQueuedProcessActionsFailedForProcess( @@ -347,24 +337,17 @@ func (s *Store) failMachineUnreachableQueuedProcessActions( } for _, row := range rows { record := processActionRecordFromSQLC(row) - if err := completeTerminalProcessActionToolCallTx( - ctx, - txNotifications, - tx, - qtx, - record, - ProcessActionStateFailed, - ProcessToolReasonMachineUnreachable, - ); err != nil { + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { return 0, err } + if record.ToolCallID != uuid.Nil { + if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return 0, err + } + } } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "machine-unreachable queued action expiry", - ); err != nil { + if err := unit.Commit(ctx, "machine-unreachable queued action expiry"); err != nil { return 0, err } return int64(len(rows)), nil @@ -410,15 +393,16 @@ func (s *Store) completeMachineUnreachableProcessActionToolCall( if err != nil { return false, err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return false, fmt.Errorf( "begin machine-unreachable process action completion: %w", err, ) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) unreachable, err := machineStillUnreachableForToolExpiryTx( ctx, @@ -432,12 +416,7 @@ func (s *Store) completeMachineUnreachableProcessActionToolCall( return false, err } if !unreachable { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "skipped machine-unreachable process action completion", - ); err != nil { + if err := unit.Commit(ctx, "skipped machine-unreachable process action completion"); err != nil { return false, err } return false, nil @@ -448,8 +427,7 @@ func (s *Store) completeMachineUnreachableProcessActionToolCall( } _, err = completeDaemonProcessActionTx( ctx, - txNotifications, - tx, + unit, qtx, processActionCompletionInput{ ProjectID: action.ProjectID, @@ -462,12 +440,7 @@ func (s *Store) completeMachineUnreachableProcessActionToolCall( state, ) if errors.Is(err, storeerr.ErrDaemonRuntimeUnregistered) { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "missed machine-unreachable process action completion", - ); err != nil { + if err := unit.Commit(ctx, "missed machine-unreachable process action completion"); err != nil { return false, err } return false, nil @@ -478,12 +451,7 @@ func (s *Store) completeMachineUnreachableProcessActionToolCall( err, ) } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "machine-unreachable process action completion", - ); err != nil { + if err := unit.Commit(ctx, "machine-unreachable process action completion"); err != nil { return false, err } return true, nil @@ -501,12 +469,12 @@ func (s *Store) completeMachineUnreachableToolCall( if err != nil { return false, err } - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return false, err } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() qtx := dbsqlc.New(tx) unreachable, err := machineStillUnreachableForToolExpiryTx( ctx, @@ -520,53 +488,45 @@ func (s *Store) completeMachineUnreachableToolCall( return false, err } if !unreachable { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "skipped machine-unreachable tool call completion", - ); err != nil { + if err := unit.Commit(ctx, "skipped machine-unreachable tool call completion"); err != nil { return false, err } return false, nil } - row, err := qtx.CompleteWaitingBuiltInToolCall( - ctx, - dbsqlc.CompleteWaitingBuiltInToolCallParams{ - ProjectID: projectID, - AgentID: agentID, - ID: toolCallID, - Outcome: string(ToolResultOutcomeFailed), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return false, nil + if _, + err := unit.LockAgent(ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: projectID, + ID: agentID}, + agentexecution.ExternalAuthority{}); err != nil { + return false, err } + process, found, err := getProcessByToolCallTx(ctx, tx, projectID, agentID, toolCallID) if err != nil { return false, err } - if _, err := finishCompletedToolCallTx( - ctx, - txNotifications, - tx, - qtx, - toolCallRecordFromWaitingCompleteSQLC(row), - toolCallResultInput{ - Outcome: ToolResultOutcomeFailed, - ResultContentParts: parts, - }, - ); err != nil { + if !found { + return false, storeerr.ErrNotFound + } + blocks, err := parseToolResultContentBlocks(parts) + if err != nil { return false, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "machine-unreachable tool call completion", - ); err != nil { + content, err := executionContent(blocks) + if err != nil { return false, err } - return true, nil + h, err := unit.Handle(projectID, agentID) + if err != nil { + return false, err + } + completed, err := h.FailProcessTool(ctx, process.ID, content) + if err != nil { + return false, err + } + if err := unit.Commit(ctx, "machine-unreachable tool call completion"); err != nil { + return false, err + } + return completed.Changed, nil } func machineStillUnreachableForToolExpiryTx( @@ -588,12 +548,15 @@ func machineStillUnreachableForToolExpiryTx( } return false, fmt.Errorf("lock machine for machine-unreachable tool expiry: %w", err) } - unreachable, err := qtx.CheckMachineUnreachableForToolExpiry(ctx, dbsqlc.CheckMachineUnreachableForToolExpiryParams{ - OrgID: orgID, - MachineID: machineID, - FallbackAt: fallbackAt, - MachineUnreachableGraceSeconds: graceSeconds, - }) + unreachable, err := qtx.CheckMachineUnreachableForToolExpiry( + ctx, + dbsqlc.CheckMachineUnreachableForToolExpiryParams{ + OrgID: orgID, + MachineID: machineID, + FallbackAt: fallbackAt, + MachineUnreachableGraceSeconds: graceSeconds, + }, + ) if errors.Is(err, pgx.ErrNoRows) { return false, nil } diff --git a/internal/storage/executionstore/process_unanswered_openings_integration_test.go b/internal/storage/executionstore/process_unanswered_openings_integration_test.go index 940005003..016722611 100644 --- a/internal/storage/executionstore/process_unanswered_openings_integration_test.go +++ b/internal/storage/executionstore/process_unanswered_openings_integration_test.go @@ -9,7 +9,10 @@ import ( "time" "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/stretchr/testify/require" ) func TestSteeringPreservesEveryUnansweredOpeningInput(t *testing.T) { @@ -126,7 +129,10 @@ func TestSteeringPreservesEveryUnansweredOpeningInput(t *testing.T) { base.Add(6*time.Second), ) if newClaim.Context.AttemptNumber != 1 { - t.Fatalf("steered context attempt = %d, want retry count reset to 1", newClaim.Context.AttemptNumber) + t.Fatalf( + "steered context attempt = %d, want retry count reset to 1", + newClaim.Context.AttemptNumber, + ) } if firstContextID != uuid.Nil { firstContext, found, err := fixture.Store.Execution().GetModelCallContext( @@ -160,14 +166,17 @@ func createUnansweredOpeningInput( now time.Time, ) executionstore.AgentInputRecord { t.Helper() - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":` + mustMarshalJSONString(t, text) + `}]`), - DeliveryMode: deliveryMode, - IdempotencyKey: idempotencyKey, - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage( + `[{"type":"text","text":` + mustMarshalJSONString(t, text) + `}]`, + ), + DeliveryMode: deliveryMode, + IdempotencyKey: idempotencyKey, + }) if err != nil { t.Fatalf("create %s: %v", idempotencyKey, err) } @@ -182,3 +191,93 @@ func mustMarshalJSONString(t *testing.T, value string) string { } return string(body) } + +func TestContinuationPreparesInheritedUnansweredOpenings(t *testing.T) { + t.Parallel() + for _, interrupted := range []string{"before_context", "retry"} { + for _, continuation := range []string{"tools", "max_tokens"} { + t.Run(interrupted+"/"+continuation, func(t *testing.T) { + t.Parallel() + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "inherited_openings_"+interrupted+"_"+continuation) + require.NoError( + t, + f.Store.Execution().ReleaseAgentRuntimeLock(ctx, testProjectID, f.AgentID, f.Lock.ID), + ) + first := ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "first") + work, found, err := f.Store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + f.Lock = work.RuntimeLock + if interrupted == "retry" { + claim := ownedModelClaim(t, ctx, f, work) + _, err = f.Store.Execution().RecordRetryableModelCallFailure(ctx, + executionstore.RecordRecoverableModelCallFailureInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, + ModelCallContextID: claim.Context.ID, ErrorKind: "transient", ErrorCode: "unavailable", + ErrorMessage: "retry later", RetryDelay: time.Hour, + }) + require.NoError(t, err) + } + require.NoError( + t, + f.Store.Execution().ReleaseAgentRuntimeLock(ctx, testProjectID, f.AgentID, f.Lock.ID), + ) + steering := ownedInput(t, ctx, f, executionstore.DeliveryModeSteering, "steering") + work, found, err = f.Store.Execution().ClaimNextAgentWork(ctx, testClaimNextAgentWorkInput()) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, []uuid.UUID{first.ID, steering.ID}, work.Model.InputIDs) + f.Lock = work.RuntimeLock + claim := ownedModelClaim(t, ctx, f, work) + if continuation == "tools" { + completeToolCallForContinuationSeedTest(t, ctx, f, work.Model.TurnID, + claim.Context.ID, "inherited", f.Now) + } else { + slug := modelProviderSlugForContext(t, ctx, f.Store, testProjectID, f.AgentID, claim.Context.ID) + _, err = f.Store.Execution().RecordModelOutputAndCompleteContext(ctx, + executionstore.RecordModelOutputAndCompleteContextInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, + ModelCallContextID: claim.Context.ID, + ProviderResponse: modelenvelope.ResponseEnvelope{ + RequestedProviderModelSlug: slug, ServedProviderModelSlug: slug, + APIFormat: modelprotocol.APIFormatOpenAIResponses, APIVariant: "default", + Normalized: modelenvelope.ResponseNormalized{ + ID: "inherited_output", StopReason: modelenvelope.StopReasonMaxTokens, + }, + }, + }) + require.NoError(t, err) + } + seed, found, err := f.Store.Execution().NextAgentModelWork(ctx, testProjectID, f.AgentID) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, executionstore.ModelWorkContinue, seed.Kind) + require.Equal(t, work.Model.InputIDs, seed.InputIDs) + prepare := executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, AgentID: f.AgentID, RuntimeLockID: f.Lock.ID, + OpeningInputIDs: seed.InputIDs, SourceModelCallContextID: seed.ModelCallContextID, + SourceModelOutputID: seed.SourceModelOutputID, + } + prepared, err := f.Store.Execution().PrepareNormalModelCall(ctx, prepare) + require.NoError(t, err) + require.True(t, prepared.Claim.Claimed) + require.Greater( + t, + prepared.Claim.Context.InputEventSequence, + claim.Context.InputEventSequence, + ) + var ids []uuid.UUID + require.NoError( + t, + f.Store.pool.QueryRow(ctx, + `SELECT opening_input_ids FROM model_call_contexts WHERE id=$1`, + prepared.Claim.Context.ID). + Scan(&ids), + ) + require.Equal(t, []uuid.UUID{first.ID, steering.ID}, ids) + + }) + } + } +} diff --git a/internal/storage/executionstore/process_wakeup_input_admission_integration_test.go b/internal/storage/executionstore/process_wakeup_input_admission_integration_test.go index 84f276498..bc2176939 100644 --- a/internal/storage/executionstore/process_wakeup_input_admission_integration_test.go +++ b/internal/storage/executionstore/process_wakeup_input_admission_integration_test.go @@ -14,7 +14,6 @@ import ( "github.com/omnara-ai/omnara/internal/interactionform" "github.com/omnara-ai/omnara/internal/storage" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/toolpermission" ) @@ -31,15 +30,20 @@ func TestConfigOnlyTurnDoesNotBecomeContinuationSeed(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - if seed, found, err := fixture.Store.Execution().NextAgentModelWork(ctx, testProjectID, fixture.AgentID); err != nil { + if seed, + found, + err := fixture.Store.Execution().NextAgentModelWork(ctx, + testProjectID, + fixture.AgentID); err != nil { t.Fatalf("load continuation seed for config-only turn: %v", err) } else if found { t.Fatalf("config-only continuation seed = %+v, want none", seed) } - if turn, err := fixture.Store.q.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: testProjectID, AgentID: fixture.AgentID}, - ); err == nil { + if turn, + err := currentContinuableTurnForTest(ctx, + fixture.Store, + testProjectID, + fixture.AgentID); err == nil { t.Fatalf("config-only current continuable turn = %+v, want none", turn) } else if !errors.Is( err, @@ -51,13 +55,11 @@ func TestConfigOnlyTurnDoesNotBecomeContinuationSeed(t *testing.T) { if wakeups := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); wakeups != 0 { t.Fatalf("config-only wakeups = %d, want 0", wakeups) } - if err := fixture.Store.Execution().MarkAgentWakeup( - ctx, - testProjectID, - fixture.AgentID, - json.RawMessage(`{"reason":"test"}`), - ); err != nil { - t.Fatalf("mark config-only wakeup: %v", err) + if _, + err := fixture.Store.pool.Exec(ctx, + `INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) VALUES($1,statement_timestamp(),statement_timestamp())`, + fixture.AgentID); err != nil { + t.Fatal(err) } claim, handled, err := fixture.Store.Execution().ClaimNextAgentWork( ctx, @@ -67,7 +69,11 @@ func TestConfigOnlyTurnDoesNotBecomeContinuationSeed(t *testing.T) { t.Fatalf("reconcile stale config-only wakeup: %v", err) } if !handled || claim.Kind != executionstore.AgentWorkNone { - t.Fatalf("stale config-only wakeup claim = %+v handled=%v, want non-executable cleanup", claim, handled) + t.Fatalf( + "stale config-only wakeup claim = %+v handled=%v, want non-executable cleanup", + claim, + handled, + ) } if wakeups := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); wakeups != 0 { t.Fatalf("config-only wakeups after claim reconciliation = %d, want 0", wakeups) @@ -86,13 +92,14 @@ func TestClaimNormalModelCallRejectsInputOutsideTurnOpening(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"valid turn input"}]`), - IdempotencyKey: "model-context-valid-turn-input", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"valid turn input"}]`), + IdempotencyKey: "model-context-valid-turn-input", + }) if err != nil { t.Fatalf("create content input: %v", err) } @@ -100,7 +107,8 @@ func TestClaimNormalModelCallRejectsInputOutsideTurnOpening(t *testing.T) { if err != nil { t.Fatalf("claim input work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) { + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) { t.Fatalf("claim = %+v found=%v, want executable input %s", claim, found, input.ID) } var configInputID uuid.UUID @@ -117,17 +125,15 @@ func TestClaimNormalModelCallRejectsInputOutsideTurnOpening(t *testing.T) { LIMIT 1`, testProjectID, fixture.AgentID).Scan(&configInputID, &configEventSequence); err != nil { t.Fatalf("load config-change input: %v", err) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent for model context: %v", err) } - _, err = fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: claim.RuntimeLock.ID, - OpeningInputIDs: []uuid.UUID{configInputID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: configEventSequence, + _, err = fixture.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: claim.RuntimeLock.ID, + OpeningInputIDs: []uuid.UUID{configInputID}, }) if !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { t.Fatalf("claim model context with non-opening source error = %v, want ErrAgentNotAdvanceable", err) @@ -146,14 +152,15 @@ func TestClaimNormalModelCallRequiresExactOpeningInputSet(t *testing.T) { ); err != nil { t.Fatalf("release fixture runtime lock: %v", err) } - first, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"first steering"}]`), - DeliveryMode: "steering", - IdempotencyKey: "exact-opening-input-first", - }) + first, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"first steering"}]`), + DeliveryMode: "steering", + IdempotencyKey: "exact-opening-input-first", + }) if err != nil { t.Fatalf("create first steering input: %v", err) } @@ -182,7 +189,7 @@ func TestClaimNormalModelCallRequiresExactOpeningInputSet(t *testing.T) { if claim.Model.InputIDs[0] != inputIDs[0] || claim.Model.InputIDs[1] != inputIDs[1] { t.Fatalf("claim input ids = %v, want %v", claim.Model.InputIDs, inputIDs) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent for model context: %v", err) } @@ -195,36 +202,32 @@ func TestClaimNormalModelCallRequiresExactOpeningInputSet(t *testing.T) { Scan(&watermark); err != nil { t.Fatalf("load opening watermark: %v", err) } - _, err = fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: claim.RuntimeLock.ID, - OpeningInputIDs: []uuid.UUID{first.ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: watermark, + _, err = fixture.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: claim.RuntimeLock.ID, + OpeningInputIDs: []uuid.UUID{first.ID}, }) if !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { t.Fatalf("missing opening input error = %v, want ErrAgentNotAdvanceable", err) } - _, err = fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: claim.RuntimeLock.ID, - OpeningInputIDs: []uuid.UUID{first.ID, first.ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: watermark, + _, err = fixture.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: claim.RuntimeLock.ID, + OpeningInputIDs: []uuid.UUID{first.ID, first.ID}, }) if !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { t.Fatalf("duplicate opening input error = %v, want ErrAgentNotAdvanceable", err) } - if _, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: claim.RuntimeLock.ID, - OpeningInputIDs: inputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: watermark, - }); err != nil { + if _, + err := fixture.Store.Execution().PrepareNormalModelCall(ctx, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: claim.RuntimeLock.ID, + OpeningInputIDs: inputIDs, + }); err != nil { t.Fatalf("claim model context with exact opening sources: %v", err) } } @@ -233,7 +236,13 @@ func TestCompletedToolResultWakeupClearsAfterLaterModelOutput(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "completed_tool_result_frontier") - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "completed_tool_result_frontier", "read_process") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "completed_tool_result_frontier", + "read_process", + ) turnID := turnIDForProcessToolCallTest(t, ctx, fixture, toolCallID) agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { @@ -322,10 +331,11 @@ func TestCompletedToolResultWakeupClearsAfterLaterModelOutput(t *testing.T) { ); err != nil { t.Fatalf("release runtime for continuation context: %v", err) } - if turn, err := fixture.Store.q.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: testProjectID, AgentID: fixture.AgentID}, - ); err == nil { + if turn, + err := currentContinuableTurnForTest(ctx, + fixture.Store, + testProjectID, + fixture.AgentID); err == nil { t.Fatalf("current continuable turn after continuation output = %+v, want none", turn) } else if !errors.Is( err, @@ -378,7 +388,8 @@ func TestResolvedInteractionCreatesWakeup(t *testing.T) { fixture := newProcessDaemonFixture(t, ctx, "resolved_interaction_wakeup") toolCallID := createToolCallForProcessTest(t, ctx, fixture, "resolved_interaction_wakeup", "ask_question") interaction := createQuestionInteractionForTest(t, ctx, fixture, toolCallID) - questionToolCall, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, toolCallID) + questionToolCall, err := fixture.Store.Execution(). + GetToolCall(ctx, testProjectID, fixture.AgentID, toolCallID) if err != nil { t.Fatalf("get question tool call: %v", err) } @@ -410,7 +421,8 @@ func TestResolvedInteractionCreatesWakeup(t *testing.T) { ); err != nil { t.Fatalf("resolve question interaction: %v", err) } - resolvedToolCall, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, toolCallID) + resolvedToolCall, err := fixture.Store.Execution(). + GetToolCall(ctx, testProjectID, fixture.AgentID, toolCallID) if err != nil { t.Fatalf("get resolved question tool call: %v", err) } @@ -431,7 +443,13 @@ func TestClaimNextAgentWorkStartsNewTurnAfterCanceledInteraction(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "claim_ignores_stopped_open_interaction") - toolCallID := createToolCallForProcessTest(t, ctx, fixture, "claim_ignores_stopped_open_interaction", "ask_question") + toolCallID := createToolCallForProcessTest( + t, + ctx, + fixture, + "claim_ignores_stopped_open_interaction", + "ask_question", + ) interaction := createQuestionInteractionForTest(t, ctx, fixture, toolCallID) cancelResult, err := fixture.Store.Execution().CancelAgent( ctx, @@ -464,10 +482,11 @@ func TestClaimNextAgentWorkStartsNewTurnAfterCanceledInteraction(t *testing.T) { if err != nil || !found || closedInteraction.State != executionstore.AgentInteractionStateCanceled { t.Fatalf("canceled interaction = %+v found=%v err=%v", closedInteraction, found, err) } - if turn, err := fixture.Store.q.CurrentContinuableAgentTurn( - ctx, - dbsqlc.CurrentContinuableAgentTurnParams{ProjectID: testProjectID, AgentID: fixture.AgentID}, - ); err == nil { + if turn, + err := currentContinuableTurnForTest(ctx, + fixture.Store, + testProjectID, + fixture.AgentID); err == nil { t.Fatalf("canceled interaction turn remained continuable: %+v", turn) } else if !errors.Is(err, pgx.ErrNoRows) { t.Fatalf("load current continuable turn after cancellation: %v", err) @@ -517,14 +536,15 @@ WHERE input.project_id = $1 if responseEvents != 0 { t.Fatalf("stopped interaction response events = %d, want 0", responseEvents) } - input, _, _, err := fixture.Store.Execution().CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - Actor: mustOmnaraActorParams(t, fixture.UserID), - ContentBlocks: json.RawMessage(`[{"type":"text","text":"next queued input"}]`), - DeliveryMode: "queued", - IdempotencyKey: "claim-ignores-stopped-open-interaction-input", - }) + input, _, _, err := fixture.Store.Execution(). + CreateAgentContentInput(ctx, executionstore.CreateAgentContentInputInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + Actor: mustOmnaraActorParams(t, fixture.UserID), + ContentBlocks: json.RawMessage(`[{"type":"text","text":"next queued input"}]`), + DeliveryMode: "queued", + IdempotencyKey: "claim-ignores-stopped-open-interaction-input", + }) if err != nil { t.Fatalf("create next queued input: %v", err) } @@ -532,7 +552,8 @@ WHERE input.project_id = $1 if err != nil { t.Fatalf("claim next work: %v", err) } - if !found || claim.Kind != executionstore.AgentWorkModel || !claimedOpeningInputIDsEqual(claim, input.ID) || + if !found || claim.Kind != executionstore.AgentWorkModel || + !claimedOpeningInputIDsEqual(claim, input.ID) || claim.Model.TurnID == interaction.TurnID { t.Fatalf( "claim found=%v claim=%+v, want new executable turn for input %s after stopped interaction turn %s", @@ -576,15 +597,17 @@ func TestPermissionApprovalReturnsToolCallToPending(t *testing.T) { if waiting.State != executionstore.ToolCallStateAwaitingPermission { t.Fatalf("tool execution after permission prompt state=%q, want awaiting_permission", waiting.State) } - if _, err := fixture.Store.Execution().ResolveAgentInteraction(ctx, executionstore.ResolveAgentInteractionInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - ID: interaction.ID, - Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{ - OptionIndices: []int{toolpermission.AllowOptionIndex}, - }}}, - Actor: mustOmnaraActorParams(t, fixture.UserID), - }); err != nil { + if _, + err := fixture.Store.Execution().ResolveAgentInteraction(ctx, + executionstore.ResolveAgentInteractionInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + ID: interaction.ID, + Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{ + OptionIndices: []int{toolpermission.AllowOptionIndex}, + }}}, + Actor: mustOmnaraActorParams(t, fixture.UserID), + }); err != nil { t.Fatalf("resolve permission interaction: %v", err) } allowed, err := fixture.Store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, toolCallID) @@ -594,6 +617,15 @@ func TestPermissionApprovalReturnsToolCallToPending(t *testing.T) { if allowed.State != executionstore.ToolCallStateReady { t.Fatalf("tool execution after permission allow state=%q, want ready", allowed.State) } + if wakeups := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); wakeups != 0 { + t.Fatalf("owned agent has %d wakeups", wakeups) + } + if err := fixture.Store.Execution().ReleaseAgentRuntimeLock(ctx, + testProjectID, + fixture.AgentID, + fixture.Lock.ID); err != nil { + t.Fatal(err) + } if wakeups := countAgentWakeups(t, ctx, fixture.Store, fixture.AgentID); wakeups != 1 { t.Fatalf("permission approval wakeups = %d, want 1", wakeups) } @@ -603,7 +635,10 @@ func TestPermissionApprovalReturnsToolCallToPending(t *testing.T) { states[0] != string(executionstore.ToolCallStateAwaitingAuthorization) || states[1] != string(executionstore.ToolCallStateAwaitingPermission) || states[2] != string(executionstore.ToolCallStateReady) { - t.Fatalf("permission tool call update states = %v, want authorization, permission, then ready", states) + t.Fatalf( + "permission tool call update states = %v, want authorization, permission, then ready", + states, + ) } } @@ -698,11 +733,10 @@ func TestCancelDuringRetryBackoffAdvancesQueuedBacklogWakeup(t *testing.T) { } } -func TestClaimNormalModelCallRejectsOpeningInputsPastWatermark(t *testing.T) { +func TestPreparationCapturesCompleteOpeningWatermark(t *testing.T) { t.Parallel() ctx := context.Background() fixture := newProcessDaemonFixture(t, ctx, "model_context_opening_watermark") - now := fixture.Now firstInput, _, _, err := fixture.Store.Execution().CreateAgentContentInput( ctx, executionstore.CreateAgentContentInputInput{ @@ -764,51 +798,53 @@ func TestClaimNormalModelCallRejectsOpeningInputsPastWatermark(t *testing.T) { admitted.Events[1].At, ) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent: %v", err) } inputIDs := []uuid.UUID{firstInput.ID, secondInput.ID} claimContext := func( openingInputIDs []uuid.UUID, - frontier int64, - at time.Time, ) (executionstore.ModelCallClaim, error) { - return fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: frontier, - }) + prepared, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, + }) + return prepared.Claim, err } if _, err := claimContext( []uuid.UUID{firstInput.ID, firstInput.ID}, - admitted.Events[1].Sequence, - now.Add(3500*time.Millisecond), ); !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("claim model context with duplicate opening input err=%v, want %v", err, storeerr.ErrAgentNotAdvanceable) + t.Fatalf( + "claim model context with duplicate opening input err=%v, want %v", + err, + storeerr.ErrAgentNotAdvanceable, + ) } if _, err := claimContext( []uuid.UUID{firstInput.ID}, - admitted.Events[1].Sequence, - now.Add(3600*time.Millisecond), - ); !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("claim model context with subset opening input err=%v, want %v", err, storeerr.ErrAgentNotAdvanceable) - } - if _, err := claimContext( - inputIDs, - admitted.Events[0].Sequence, - now.Add(4*time.Second), ); !errors.Is(err, storeerr.ErrAgentNotAdvanceable) { - t.Fatalf("claim model context before second opening event err=%v, want %v", err, storeerr.ErrAgentNotAdvanceable) + t.Fatalf( + "claim model context with subset opening input err=%v, want %v", + err, + storeerr.ErrAgentNotAdvanceable, + ) } - firstClaim, err := claimContext(inputIDs, admitted.Events[1].Sequence, now.Add(5*time.Second)) + firstClaim, err := claimContext(inputIDs) if err != nil { t.Fatalf("claim model context after full opening set: %v", err) } - repeatedClaim, err := claimContext(inputIDs, admitted.Events[1].Sequence, now.Add(6*time.Second)) + if firstClaim.Context.InputEventSequence != admitted.Events[1].Sequence { + t.Fatalf( + "prepared watermark=%d want=%d", + firstClaim.Context.InputEventSequence, + admitted.Events[1].Sequence, + ) + } + repeatedClaim, err := claimContext(inputIDs) if err != nil { t.Fatalf("repeat semantic model context claim: %v", err) } diff --git a/internal/storage/executionstore/processes_store.go b/internal/storage/executionstore/processes_store.go index df88aa6e0..679773045 100644 --- a/internal/storage/executionstore/processes_store.go +++ b/internal/storage/executionstore/processes_store.go @@ -10,6 +10,7 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/daemonprotocol" "github.com/omnara-ai/omnara/internal/processcmd" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -34,7 +35,9 @@ func (s *Store) GetDaemonGitCredentialsScope( orgID, machineID, processID uuid.UUID, ) (DaemonGitCredentialsScope, bool, error) { if orgID == uuid.Nil || machineID == uuid.Nil || processID == uuid.Nil { - return DaemonGitCredentialsScope{}, false, errors.New("organization, machine, and process are required") + return DaemonGitCredentialsScope{}, false, errors.New( + "organization, machine, and process are required", + ) } row, err := s.q.GetDaemonGitCredentialsScope(ctx, dbsqlc.GetDaemonGitCredentialsScopeParams{ OrgID: orgID, MachineID: machineID, ProcessID: processID, @@ -119,6 +122,7 @@ func (t *toolCallTransaction) startProcess( return ProcessRecord{}, storeerr.ErrIdempotencyConflict } t.hasDurableCompletionOwner = true + t.owner.ProcessID = existing.ID if err := t.lockOrAcceptExisting(ctx); err != nil { return ProcessRecord{}, err } @@ -229,6 +233,8 @@ func (t *toolCallTransaction) startProcess( if !processReplayMatches(existing, input) { return ProcessRecord{}, storeerr.ErrIdempotencyConflict } + t.owner.ProcessID = existing.ID + t.hasDurableCompletionOwner = true return existing, nil } if _, found, loadErr := machineReachableForProjectMachineTx( @@ -246,6 +252,7 @@ func (t *toolCallTransaction) startProcess( return ProcessRecord{}, fmt.Errorf("start process: %w", err) } process := processRecordFromInsertSQLC(row) + t.owner.ProcessID = process.ID t.hasDurableCompletionOwner = true t.requiresWaitingDisposition = true t.notifications.AddDaemonWork(process.MachineID) @@ -322,16 +329,18 @@ func (s *Store) CompleteProcess(ctx context.Context, input CompleteProcessInput) return ProcessRecord{}, errors.New("source ended at is required") } input.SourceEndedAt = canonicalSourceTime(input.SourceEndedAt) - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return ProcessRecord{}, fmt.Errorf("begin complete process: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) - if _, err := qtx.LockAgentInProject( + if _, err := unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, + agentexecution.ExternalAuthority{}, ); err != nil { return ProcessRecord{}, fmt.Errorf("lock agent for process completion: %w", err) } @@ -356,14 +365,13 @@ func (s *Store) CompleteProcess(ctx context.Context, input CompleteProcessInput) process := processRecordFromCompleteSQLC(row) if err := completeQueuedProcessActionsForTerminalProcessTx( ctx, - txNotifications, - tx, + unit, qtx, process, ); err != nil { return ProcessRecord{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "complete process"); err != nil { + if err := unit.Commit(ctx, "complete process"); err != nil { return ProcessRecord{}, err } return process, nil diff --git a/internal/storage/executionstore/processes_test_helpers_test.go b/internal/storage/executionstore/processes_test_helpers_test.go index d309c8bd5..b6619a3ef 100644 --- a/internal/storage/executionstore/processes_test_helpers_test.go +++ b/internal/storage/executionstore/processes_test_helpers_test.go @@ -28,6 +28,7 @@ import ( "github.com/omnara-ai/omnara/internal/testutil/storagetest" "github.com/omnara-ai/omnara/internal/toolcatalog" "github.com/omnara-ai/omnara/internal/toolpermission" + "github.com/stretchr/testify/require" ) func createQuestionInteractionForTest( @@ -137,17 +138,24 @@ func insertToolCallForTest( return executionstore.ToolCallRecord{}, err } defer func() { _ = tx.Rollback(ctx) }() - row, err := store.q.WithTx(tx).InsertToolCall(ctx, dbsqlc.InsertToolCallParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - SourceEventID: input.SourceEventID, - ModelCallContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ProviderCallID: input.ProviderCallID, - Name: input.Name, - Input: input.Input, - Type: input.Type, - }) + var callID uuid.UUID + err = tx.QueryRow(ctx, ` +INSERT INTO tool_calls(agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) +SELECT event.agent_id,output.id,$6,$7,$8,$9,'awaiting_authorization',output.created_at +FROM agent_events event +JOIN agents agent ON agent.id=event.agent_id +JOIN model_outputs output ON output.agent_id=event.agent_id AND output.id=event.model_output_id +JOIN model_call_contexts context ON context.agent_id=output.agent_id AND context.id=output.model_call_context_id +JOIN agent_runtime_locks runtime ON runtime.agent_id=context.agent_id AND runtime.id=context.runtime_lock_id +WHERE agent.project_id=$1 AND event.agent_id=$2 AND event.id=$3 AND context.id=$4 + AND runtime.id=$5 AND runtime.cancel_requested_at IS NULL AND runtime.lease_expires_at>statement_timestamp() +RETURNING id`, + input.ProjectID, + input.AgentID, + input.SourceEventID, + input.ModelCallContextID, + input.RuntimeLockID, + input.ProviderCallID, input.Name, input.Input, input.Type).Scan(&callID) if storeutil.IsUniqueViolation(err) { return executionstore.ToolCallRecord{}, storeerr.ErrIdempotencyConflict } @@ -179,7 +187,7 @@ WHERE agent.project_id = $1 AND call.agent_id = $2 AND call.id = $3 GROUP BY call.model_output_id -`, input.ProjectID, input.AgentID, row.ID).Scan(&modelOutputID, &ordinal); err != nil { +`, input.ProjectID, input.AgentID, callID).Scan(&modelOutputID, &ordinal); err != nil { return executionstore.ToolCallRecord{}, err } if _, err := executionstore.IntegrationCreateContentBlockTx(ctx, tx, executionstore.CreateContentBlockInput{ @@ -189,14 +197,17 @@ GROUP BY call.model_output_id OwnerModelOutputID: modelOutputID, Ordinal: ordinal, BlockKind: executionstore.ContentBlockKindToolCall, - ToolCallID: row.ID, + ToolCallID: callID, }); err != nil { return executionstore.ToolCallRecord{}, err } if err := tx.Commit(ctx); err != nil { return executionstore.ToolCallRecord{}, err } - return executionstore.IntegrationToolCallRecordFromInsertSQLC(row), nil + if err := store.Execution().RepairExecutionForTest(ctx, input.ProjectID, input.AgentID); err != nil { + return executionstore.ToolCallRecord{}, err + } + return store.Execution().GetToolCall(ctx, input.ProjectID, input.AgentID, callID) } type recordingPostCommitPublisher struct { @@ -231,7 +242,10 @@ func createProcessActionForTest( ) } -func (p *recordingPostCommitPublisher) PublishPostCommit(_ context.Context, intent notifications.PostCommitIntent) { +func (p *recordingPostCommitPublisher) PublishPostCommit( + _ context.Context, + intent notifications.PostCommitIntent, +) { p.intents = append(p.intents, intent) } @@ -383,7 +397,9 @@ func (fixture processDaemonFixture) authority() executionstore.DaemonRuntimeAuth } } -func (fixture processDaemonFixture) authorityForRuntime(runtimeID uuid.UUID) executionstore.DaemonRuntimeAuthority { +func (fixture processDaemonFixture) authorityForRuntime( + runtimeID uuid.UUID, +) executionstore.DaemonRuntimeAuthority { authority := fixture.authority() authority.DaemonRuntimeID = runtimeID return authority @@ -393,7 +409,10 @@ func newProcessDaemonFixture(t *testing.T, ctx context.Context, testName string) t.Helper() pool := openIntegrationDB(t, ctx) seedMigratedDB(t, ctx, pool) - store := newIntegrationStore(pool, storage.WithModelCallRetryBackoff(func(int, string) time.Duration { return 0 })) + store := newIntegrationStore( + pool, + storage.WithModelCallRetryBackoff(func(int, string) time.Duration { return 0 }), + ) now := time.Date(2026, 5, 17, 10, 0, 0, 0, time.UTC) user := mustCreateProjectOperatorUser(t, ctx, store, "process-"+testName+"@example.com", "Process Tester") return newProcessDaemonFixtureInStore(t, ctx, store, user.ID, testName, now) @@ -793,6 +812,48 @@ func createToolCallBatchForProcessTest( if len(items) == 0 { t.Fatal("process tool-call batch must contain at least one proposal") } + var sourceID, contextID uuid.UUID + err := fixture.Store.pool.QueryRow(ctx, + `SELECT event.id,output.model_call_context_id FROM agent_execution_state head JOIN model_outputs output ON output.agent_id=head.agent_id AND output.id=head.pending_tool_output_id JOIN agent_events event ON event.agent_id=output.agent_id AND event.model_output_id=output.id WHERE head.agent_id=$1 AND head.incomplete_tools`, + fixture.AgentID). + Scan(&sourceID, &contextID) + if err == nil { + ids := make([]uuid.UUID, 0, len(items)) + for _, item := range items { + raw := item.Input + if len(raw) == 0 { + raw = json.RawMessage(`{}`) + } + call, err := insertToolCallForTest( + ctx, + fixture.Store, + toolCallFixtureInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + SourceEventID: sourceID, + ModelCallContextID: contextID, + RuntimeLockID: fixture.Lock.ID, + ProviderCallID: "call_" + item.TestName, + Name: item.ToolName, + Input: raw, + Type: item.ToolType, + }, + ) + require.NoError(t, err) + if item.Allowed { + _, err = fixture.Store.Execution(). + MarkToolCallReady(ctx, + executionstore.MarkToolCallReadyInput{ProjectID: testProjectID, + AgentID: fixture.AgentID, + ID: call.ID, + RuntimeLockID: fixture.Lock.ID}) + require.NoError(t, err) + } + ids = append(ids, call.ID) + } + return ids + } + require.ErrorIs(t, err, pgx.ErrNoRows) input, _, _, err := fixture.Store.Execution().CreateAgentContentInput( ctx, executionstore.CreateAgentContentInputInput{ @@ -801,6 +862,7 @@ func createToolCallBatchForProcessTest( Actor: mustOmnaraActorParams(t, fixture.UserID), ContentBlocks: json.RawMessage(`[{"type":"text","text":"seed tool call"}]`), IdempotencyKey: "process-tool-input-" + batchName, + DeliveryMode: executionstore.DeliveryModeSteering, }, ) if err != nil { @@ -817,18 +879,19 @@ func createToolCallBatchForProcessTest( if !found || len(admitted.Events) != 1 { t.Fatalf("admit source agent input: found=%v admitted=%+v", found, admitted) } - agent, err := fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) + _, err = fixture.Store.Execution().GetAgentInProject(ctx, testProjectID, fixture.AgentID) if err != nil { t.Fatalf("load agent for model context fixture: %v", err) } - claim, err := fixture.Store.Execution().ClaimNormalModelCall(ctx, executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: agent.CurrentConfigID, - InputEventSequence: admitted.Events[0].Sequence, - }) + prepared1, err := fixture.Store.Execution(). + PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, + }) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim model call context: %v", err) } @@ -1344,25 +1407,13 @@ WHERE agent.project_id = $1 ); err != nil { t.Fatalf("append forced tool result event: %v", err) } - metadata, err := marshalJSON( - map[string]any{ - "reason": "tool_result", - "tool_call_id": record.ID, - "model_call_context_id": record.ModelCallContextID, - "outcome": outcome, - }, - ) - if err != nil { - t.Fatalf("marshal forced tool result wakeup metadata: %v", err) - } - if err := dbsqlc.New(tx).MarkAgentWakeup( - ctx, dbsqlc.MarkAgentWakeupParams{ProjectID: projectID, AgentID: agentID, Metadata: metadata}, - ); err != nil { - t.Fatalf("mark forced tool result wakeup: %v", err) - } if err := tx.Commit(ctx); err != nil { t.Fatalf("commit forced tool completion: %v", err) } + if err := store.Execution().RepairExecutionForTest(ctx, projectID, agentID); err != nil { + t.Fatal(err) + } + } func cancelToolCallForTest( @@ -1441,7 +1492,11 @@ func assertCompletedToolCallWithResult( t.Fatalf("completed tool call has no result authority: %+v", toolCall) } if result.Outcome != toolCall.Outcome { - t.Fatalf("tool result outcome %q does not match execution outcome %q", result.Outcome, toolCall.Outcome) + t.Fatalf( + "tool result outcome %q does not match execution outcome %q", + result.Outcome, + toolCall.Outcome, + ) } if reason != "" { completed := completedToolCallForTest(t, store, agentID, toolCall.TurnID, toolCall.ID) diff --git a/internal/storage/executionstore/project_machine_pool_grants_store.go b/internal/storage/executionstore/project_machine_pool_grants_store.go index 9ddd648bd..b44aa50d6 100644 --- a/internal/storage/executionstore/project_machine_pool_grants_store.go +++ b/internal/storage/executionstore/project_machine_pool_grants_store.go @@ -11,6 +11,8 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/resourcemeta" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -18,8 +20,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/management" "github.com/omnara-ai/omnara/internal/storage/patch" "github.com/omnara-ai/omnara/internal/storage/storeerr" - - "github.com/omnara-ai/omnara/internal/resourcemeta" ) type ProjectMachinePoolGrantRecord struct { @@ -772,12 +772,13 @@ func (s *Store) deleteProjectMachinePoolGrantOnce( ctx context.Context, orgID, projectID, id uuid.UUID, ) (DeleteProjectMachinePoolGrantResult, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return DeleteProjectMachinePoolGrantResult{}, fmt.Errorf("begin revoke project machine pool grant: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() qtx := dbsqlc.New(tx) if err := lifecyclelock.EnterActiveProject(ctx, tx, orgID, projectID); err != nil { return DeleteProjectMachinePoolGrantResult{}, err @@ -856,7 +857,7 @@ func (s *Store) deleteProjectMachinePoolGrantOnce( AgentID: agentRow.AgentID, }) } - if err := lifecyclelock.Agents(ctx, tx, agentRefs); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { return DeleteProjectMachinePoolGrantResult{}, err } machineRows, err := qtx.MarkPoolGrantMachinesDeleting( @@ -875,7 +876,7 @@ func (s *Store) deleteProjectMachinePoolGrantOnce( if err := completeExecutionRevokedProcessesTx( ctx, txNotifications, - tx, + unit, qtx, executionRevokedProcessScope{ projectID: projectID, @@ -895,7 +896,7 @@ func (s *Store) deleteProjectMachinePoolGrantOnce( } return DeleteProjectMachinePoolGrantResult{}, fmt.Errorf("delete project machine pool grant: %w", err) } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "delete project machine pool grant"); err != nil { + if err := unit.Commit(ctx, "delete project machine pool grant"); err != nil { return DeleteProjectMachinePoolGrantResult{}, err } result := DeleteProjectMachinePoolGrantResult{ diff --git a/internal/storage/executionstore/provider_cost.go b/internal/storage/executionstore/provider_cost.go index ce9e78315..f501efbf2 100644 --- a/internal/storage/executionstore/provider_cost.go +++ b/internal/storage/executionstore/provider_cost.go @@ -4,14 +4,6 @@ import ( "github.com/omnara-ai/omnara/internal/modelenvelope" ) -func providerReportedCostUSDToSQLC(cost modelenvelope.ProviderReportedCostUSD) *string { - if cost == "" { - return nil - } - value := string(cost) - return &value -} - func providerReportedCostUSDFromSQLC(value string) modelenvelope.ProviderReportedCostUSD { if value == "" { return "" diff --git a/internal/storage/executionstore/provider_cost_test.go b/internal/storage/executionstore/provider_cost_test.go index 22a49d0cb..50ef8fd25 100644 --- a/internal/storage/executionstore/provider_cost_test.go +++ b/internal/storage/executionstore/provider_cost_test.go @@ -14,11 +14,7 @@ func TestProviderReportedCostUSDSQLCRoundTrip(t *testing.T) { "100", } { t.Run(string(cost), func(t *testing.T) { - sqlValue := providerReportedCostUSDToSQLC(cost) - stored := "" - if sqlValue != nil { - stored = *sqlValue - } + stored := string(cost) if got := providerReportedCostUSDFromSQLC(stored); got != cost { t.Fatalf("provider-reported cost round trip = %q, want %q", got, cost) } diff --git a/internal/storage/executionstore/provider_runtime_reconciliation_store.go b/internal/storage/executionstore/provider_runtime_reconciliation_store.go index 0f83c10db..5f5af1b76 100644 --- a/internal/storage/executionstore/provider_runtime_reconciliation_store.go +++ b/internal/storage/executionstore/provider_runtime_reconciliation_store.go @@ -243,7 +243,7 @@ func (s *Store) MarkProviderRuntimeMismatch( return false, fmt.Errorf("begin provider runtime mismatch mark: %w", err) } defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + qtx := dbsqlc.New(tx) if err := lifecyclelock.Machines( ctx, @@ -348,16 +348,17 @@ func (s *Store) ClaimProviderRuntimeMismatchDeletion( if !providerRuntimeWakeAllowsDeletion(candidate) { return PoolMachineDeletionClaim{}, false, nil } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return PoolMachineDeletionClaim{}, false, fmt.Errorf( "begin provider runtime mismatch deletion claim: %w", err, ) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) - txNotifications := s.newTxNotifications() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) locked, err := lockProviderRuntimeDeletionCandidate(ctx, tx, qtx, candidate) if err != nil { return PoolMachineDeletionClaim{}, false, err @@ -395,16 +396,19 @@ func (s *Store) ClaimProviderRuntimeMismatchDeletion( claim := providerRuntimeMismatchDeletionClaimFromSQLC(row) claim, err = s.finalizePoolMachineDeletionClaimTx( ctx, - tx, + unit, qtx, txNotifications, claim, - "claim provider runtime mismatch deletion", + ProcessToolReasonMachineUnreachable, ) if err != nil { return PoolMachineDeletionClaim{}, false, err } + if err := unit.Commit(ctx, "claim provider runtime mismatch deletion"); err != nil { + return PoolMachineDeletionClaim{}, false, err + } return claim, true, nil } @@ -420,15 +424,16 @@ func (s *Store) ClaimProviderRuntimeTerminatedDeletion( if !providerRuntimeWakeAllowsDeletion(candidate) { return PoolMachineDeletionClaim{}, false, nil } - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return PoolMachineDeletionClaim{}, false, fmt.Errorf( "begin provider runtime terminated deletion claim: %w", err, ) } - defer func() { _ = tx.Rollback(ctx) }() - qtx := s.q.WithTx(tx) + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) locked, err := lockProviderRuntimeDeletionCandidate(ctx, tx, qtx, candidate) if err != nil { return PoolMachineDeletionClaim{}, false, err @@ -464,16 +469,18 @@ func (s *Store) ClaimProviderRuntimeTerminatedDeletion( claim := providerRuntimeTerminatedDeletionClaimFromSQLC(row) claim, err = s.finalizePoolMachineDeletionClaimTx( ctx, - tx, + unit, qtx, - s.newTxNotifications(), + unit.Notifications(), claim, - "claim provider runtime terminated deletion", ProcessToolReasonMachineUnreachable, ) if err != nil { return PoolMachineDeletionClaim{}, false, err } + if err := unit.Commit(ctx, "claim provider runtime terminated deletion"); err != nil { + return PoolMachineDeletionClaim{}, false, err + } return claim, true, nil } @@ -489,7 +496,7 @@ func providerRuntimeWakeAllowsDeletion(candidate ProviderRuntimeCandidate) bool func lockProviderRuntimeDeletionCandidate( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, qtx *dbsqlc.Queries, candidate ProviderRuntimeCandidate, ) (bool, error) { diff --git a/internal/storage/executionstore/store.go b/internal/storage/executionstore/store.go index 9bb48c768..11e8d370d 100644 --- a/internal/storage/executionstore/store.go +++ b/internal/storage/executionstore/store.go @@ -1,15 +1,14 @@ package executionstore import ( - "context" "time" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/secretstore" @@ -26,65 +25,31 @@ type Config struct { } type Store struct { - pool *storeutil.Pool - q *dbsqlc.Queries - postCommitPublisher notifications.PostCommitPublisher - modelCallRetryBackoff func(int, string) time.Duration - integrations *integrationstore.Store - machinePoolProviders MachinePoolProviders - identity *identitystore.Store - secrets *secretstore.Store - artifacts *artifactstore.Store + cell *agentexecution.AgentCell + pool *storeutil.Pool + q *dbsqlc.Queries + integrations *integrationstore.Store + machinePoolProviders MachinePoolProviders + identity *identitystore.Store + secrets *secretstore.Store + artifacts *artifactstore.Store } func New(pool *pgxpool.Pool, config Config) *Store { db := storeutil.WrapPool(pool) return &Store{ - pool: db, - q: dbsqlc.New(db), - postCommitPublisher: config.PostCommitPublisher, - modelCallRetryBackoff: config.ModelCallRetryBackoff, - integrations: config.Integrations, - machinePoolProviders: config.MachinePoolProviders, - identity: config.Identity, - secrets: config.Secrets, - artifacts: config.Artifacts, + cell: agentexecution.NewCell( + "primary", + pool, + config.PostCommitPublisher, + config.ModelCallRetryBackoff, + ), + pool: db, + q: dbsqlc.New(db), + integrations: config.Integrations, + machinePoolProviders: config.MachinePoolProviders, + identity: config.Identity, + secrets: config.Secrets, + artifacts: config.Artifacts, } } - -func (s *Store) modelCallRetryDelay(attemptNumber int, contextID string) time.Duration { - if s.modelCallRetryBackoff == nil { - return ModelCallRetryBackoff(attemptNumber, contextID) - } - delay := s.modelCallRetryBackoff(attemptNumber, contextID) - if delay < 0 { - return 0 - } - return delay -} - -func (s *Store) newTxNotifications() *notifications.TxNotifications { - return notifications.NewTxNotifications() -} - -func (s *Store) commitTxWithNotifications( - ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, - operation string, -) error { - if err := enqueueEventWebhooksTx(ctx, tx, txNotifications); err != nil { - return err - } - var publisher notifications.PostCommitPublisher - if s != nil { - publisher = s.postCommitPublisher - } - return storeutil.CommitTxWithNotifications( - ctx, - tx, - txNotifications, - publisher, - operation, - ) -} diff --git a/internal/storage/executionstore/subagents_integration_test.go b/internal/storage/executionstore/subagents_integration_test.go index 114a62884..33e201f10 100644 --- a/internal/storage/executionstore/subagents_integration_test.go +++ b/internal/storage/executionstore/subagents_integration_test.go @@ -606,8 +606,8 @@ func TestArchiveIdleAgentsWaitsForBusyDescendants(t *testing.T) { if err := store.Execution().ReleaseAgentRuntimeLock(ctx, testProjectID, leaf.Agent.ID, leafLock.ID); err != nil { t.Fatalf("release leaf runtime lock: %v", err) } - if err := store.Execution().MarkAgentWakeup( - ctx, testProjectID, leaf.Agent.ID, []byte(`{"reason":"test"}`), + if err := store.Execution().RepairExecutionForTest( + ctx, testProjectID, leaf.Agent.ID, ); err != nil { t.Fatalf("mark leaf wakeup: %v", err) } diff --git a/internal/storage/executionstore/subagents_store.go b/internal/storage/executionstore/subagents_store.go index c46725090..7761c0236 100644 --- a/internal/storage/executionstore/subagents_store.go +++ b/internal/storage/executionstore/subagents_store.go @@ -5,15 +5,15 @@ import ( "encoding/json" "errors" "fmt" + "slices" "strings" "time" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/agentconfig" - "github.com/omnara-ai/omnara/internal/interactionform" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -56,13 +56,6 @@ type SubagentStatus struct { HasOpenPermission bool } -type subagentMessage struct { - Kind string - Text string - InteractionID uuid.UUID - IdempotencyKey string -} - func SubagentActorParams(orgID uuid.UUID, agent AgentRecord) (*ActorParams, error) { tenantID, err := publicid.Encode(publicid.KindOrganization, orgID) if err != nil { @@ -89,7 +82,7 @@ func subagentDisplayName(agent AgentRecord) string { return "agent" } -func lockSubagentParentSourcesTx(ctx context.Context, tx pgx.Tx, launch SubagentLaunch) error { +func lockSubagentParentSourcesTx(ctx context.Context, tx dbsqlc.DBTX, launch SubagentLaunch) error { if launch.ParentAgentID == uuid.Nil || launch.Key == "" { return errors.New("subagent launch requires a parent agent and key") } @@ -103,20 +96,19 @@ func (s *Store) AgentDepth(ctx context.Context, projectID, agentID uuid.UUID) (i if projectID == uuid.Nil || agentID == uuid.Nil { return 0, errors.New("project id and agent id are required") } - depth, err := s.q.CountAgentAncestors(ctx, dbsqlc.CountAgentAncestorsParams{ProjectID: projectID, AgentID: agentID}) + depth, err := s.q.CountAgentAncestors( + ctx, + dbsqlc.CountAgentAncestorsParams{ProjectID: projectID, AgentID: agentID}, + ) if err != nil { return 0, fmt.Errorf("count agent ancestors: %w", err) } return int(depth), nil } -// lockParentMachineBindingsForSharingTx lists the parent's attached machines -// and takes their lifecycle locks. It runs after the child row is inserted so -// subagent launches take the agent quota before machine locks, in the same -// order as top-level launches take the quota before pool and machine locks. func lockParentMachineBindingsForSharingTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, qtx *dbsqlc.Queries, orgID, projectID, parentAgentID uuid.UUID, ) ([]dbsqlc.ListParentMachineBindingsForSharingRow, error) { @@ -134,20 +126,42 @@ func lockParentMachineBindingsForSharingTx( if err := lifecyclelock.Machines(ctx, tx, machineRefs); err != nil { return nil, err } - return sharedBindings, nil + revalidated, err := qtx.ListParentMachineBindingsForSharing( + ctx, + dbsqlc.ListParentMachineBindingsForSharingParams{ + ProjectID: projectID, AgentID: parentAgentID, + }, + ) + if err != nil { + return nil, fmt.Errorf("revalidate parent machine bindings: %w", err) + } + if !slices.EqualFunc( + sharedBindings, + revalidated, + func(a, b dbsqlc.ListParentMachineBindingsForSharingRow) bool { + return a.MachineID == b.MachineID && a.ProjectMachineGrantID == b.ProjectMachineGrantID + }, + ) { + return nil, fmt.Errorf( + "parent machine bindings changed during launch: %w", + storeutil.ErrRetryTransaction, + ) + } + return revalidated, nil } func admitSubagentLaunchTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, projectID uuid.UUID, launch SubagentLaunch, ) error { - if err := lifecyclelock.Agents(ctx, tx, []lifecyclelock.AgentRef{{ + tx := unit.DB() + if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ ProjectID: projectID, AgentID: launch.ParentAgentID, - }}); err != nil { + }}, agentexecution.LifecycleAuthority{}); err != nil { if errors.Is(err, storeerr.ErrNotFound) { return fmt.Errorf("parent agent: %w", storeerr.ErrNotFound) } @@ -233,7 +247,12 @@ func subagentStatusFromSQLC(row dbsqlc.ListChildAgentsRow) SubagentStatus { HasOpenQuestion: row.HasOpenQuestion, HasOpenPermission: row.HasOpenPermission, } - status.State = agentActivityState(status.Archived, status.HasOpenQuestion, status.HasOpenPermission, status.IsRunning) + status.State = agentActivityState( + status.Archived, + status.HasOpenQuestion, + status.HasOpenPermission, + status.IsRunning, + ) return status } @@ -302,40 +321,24 @@ func resolveSubagentTx( return child, nil } } - return SubagentStatus{}, storeerr.InvalidRequest(fmt.Errorf("no subagent matches agent_id %q", agentPublicID)) -} - -func renderQuestionForParent(form interactionform.Form) string { - var builder strings.Builder - builder.WriteString("Question: ") - builder.WriteString(form.Title) - for _, item := range form.Context { - builder.WriteString("\n") - builder.WriteString(item.Label) - builder.WriteString(": ") - builder.WriteString(item.Value) - } - for index, question := range form.Questions { - builder.WriteString(fmt.Sprintf("\n%d. %s", index+1, question.Prompt)) - for _, option := range question.Options { - builder.WriteString("\n - ") - builder.WriteString(option.Label) - } - } - return builder.String() + return SubagentStatus{}, storeerr.InvalidRequest( + fmt.Errorf("no subagent matches agent_id %q", agentPublicID), + ) } -func lockAgentWithParentTx(ctx context.Context, tx pgx.Tx, qtx *dbsqlc.Queries, projectID, agentID uuid.UUID) error { +func lockAgentWithParentTx( + ctx context.Context, unit *agentexecution.Unit, qtx *dbsqlc.Queries, projectID, agentID uuid.UUID, +) error { refs, err := agentWithParentLockRefsTx(ctx, qtx, projectID, agentID) if err != nil { return err } - return lifecyclelock.Agents(ctx, tx, refs) + return unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}) } func tryLockAgentWithParentTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, qtx *dbsqlc.Queries, projectID, agentID uuid.UUID, ) (bool, error) { @@ -343,7 +346,11 @@ func tryLockAgentWithParentTx( if err != nil { return false, err } - return lifecyclelock.TryAgents(ctx, tx, refs) + plan, err := unit.PlanAgents(ctx, refs, agentexecution.LifecycleAuthority{}) + if err != nil { + return false, err + } + return unit.TryLockAgents(ctx, plan) } func agentWithParentLockRefsTx( @@ -351,7 +358,10 @@ func agentWithParentLockRefsTx( qtx *dbsqlc.Queries, projectID, agentID uuid.UUID, ) ([]lifecyclelock.AgentRef, error) { - parentID, err := qtx.GetAgentParentID(ctx, dbsqlc.GetAgentParentIDParams{ProjectID: projectID, ID: agentID}) + parentID, err := qtx.GetAgentParentID( + ctx, + dbsqlc.GetAgentParentIDParams{ProjectID: projectID, ID: agentID}, + ) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return nil, storeerr.ErrNotFound @@ -365,120 +375,6 @@ func agentWithParentLockRefsTx( return refs, nil } -func notifyParentAgentTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - child AgentRecord, - message subagentMessage, -) error { - parent, err := loadAgentInProjectTx(ctx, tx, child.ProjectID, child.ParentAgentID) - if err != nil { - return err - } - if parent.State == AgentStateArchived { - return nil - } - actor, err := SubagentActorParams(child.OrgID, child) - if err != nil { - return err - } - childPublicID, err := publicid.Encode(publicid.KindAgent, child.ID) - if err != nil { - return fmt.Errorf("encode subagent id: %w", err) - } - metadataBody := map[string]any{ - "kind": message.Kind, - "agent_id": childPublicID, - "name": child.Name, - "key": child.SubagentKey, - } - if message.InteractionID != uuid.Nil { - interactionPublicID, err := publicid.Encode(publicid.KindAgentInteraction, message.InteractionID) - if err != nil { - return fmt.Errorf("encode interaction id: %w", err) - } - metadataBody["interaction_id"] = interactionPublicID - } - metadata, err := marshalJSON(map[string]any{"subagent_message": metadataBody}) - if err != nil { - return fmt.Errorf("marshal subagent message metadata: %w", err) - } - contentBlocks, contentBlocksJSON, err := textInputContentBlocks(subagentMessageText(child, childPublicID, message)) - if err != nil { - return err - } - if _, err := createAgentContentInputTx(ctx, txNotifications, tx, qtx, CreateAgentContentInputInput{ - ProjectID: parent.ProjectID, - AgentID: parent.ID, - Actor: actor, - ContentBlocks: contentBlocksJSON, - Metadata: metadata, - DeliveryMode: DeliveryModeSteering, - IdempotencyScope: subagentMessageIdempotencyScope, - IdempotencyKey: message.IdempotencyKey, - }, contentBlocks); err != nil { - if errors.Is(err, storeerr.ErrStateTransitionConflict) { - return nil - } - return fmt.Errorf("deliver subagent message to parent: %w", err) - } - return nil -} - -func subagentMessageText(child AgentRecord, childPublicID string, message subagentMessage) string { - label := fmt.Sprintf( - "Subagent %q (agent_id %s, key %q)", subagentDisplayName(child), childPublicID, child.SubagentKey, - ) - var header string - switch message.Kind { - case SubagentMessageKindResult: - header = label + " finished its turn:" - case SubagentMessageKindRefused: - header = label + " stopped because its model refused to continue; it produced no final answer:" - case SubagentMessageKindContentFiltered: - header = label + " stopped because its model output was blocked by a content filter; " + - "it produced no final answer:" - case SubagentMessageKindFailed: - header = label + " failed:" - case SubagentMessageKindQuestion: - header = label + " asked a question and is paused until a human answers it. " + - "Messaging it with send_agent_message cancels the question." - default: - header = label + ":" - } - if strings.TrimSpace(message.Text) == "" { - return header - } - return header + "\n\n" + message.Text -} - -func handleSubagentTurnEndedTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, - message subagentMessage, -) error { - parentID, err := qtx.GetAgentParentID(ctx, dbsqlc.GetAgentParentIDParams{ProjectID: projectID, ID: agentID}) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrNotFound - } - return fmt.Errorf("load agent parent: %w", err) - } - if parentID == nil { - return nil - } - child, err := loadAgentInProjectTx(ctx, tx, projectID, agentID) - if err != nil { - return err - } - return notifyParentAgentTx(ctx, txNotifications, tx, qtx, child, message) -} - func textInputContentBlocks(text string) ([]CreateContentBlockInput, json.RawMessage, error) { contentBlocks := []CreateContentBlockInput{{ Ordinal: 0, @@ -492,32 +388,6 @@ func textInputContentBlocks(text string) ([]CreateContentBlockInput, json.RawMes return contentBlocks, contentBlocksJSON, nil } -func handleSubagentQuestionTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - interaction AgentInteractionRecord, -) error { - child, err := loadAgentInProjectTx(ctx, tx, interaction.ProjectID, interaction.AgentID) - if err != nil { - return err - } - if child.ParentAgentID == uuid.Nil { - return nil - } - form, err := interaction.Form() - if err != nil { - return err - } - return notifyParentAgentTx(ctx, txNotifications, tx, qtx, child, subagentMessage{ - Kind: SubagentMessageKindQuestion, - Text: renderQuestionForParent(form), - InteractionID: interaction.ID, - IdempotencyKey: "question:" + interaction.ID.String(), - }) -} - type ListAgentInteractionsInput struct { ProjectID uuid.UUID AgentIDs []uuid.UUID @@ -617,10 +487,12 @@ func SendSubagentMessageForToolCall( } func (t *toolCallTransaction) sendSubagentMessage(ctx context.Context, input SendSubagentMessageInput) error { - if err := lifecyclelock.Agents(ctx, t.tx, []lifecyclelock.AgentRef{ + if err := t.unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{ {ProjectID: t.input.ProjectID, AgentID: t.input.AgentID}, {ProjectID: t.input.ProjectID, AgentID: input.TargetAgentID}, - }); err != nil { + }, + agentexecution.RuntimeAuthority{AgentID: t.input.AgentID, + RuntimeLockID: t.input.RuntimeLockID}); err != nil { return err } if err := t.lockForMutation(ctx); err != nil { @@ -656,7 +528,7 @@ func (t *toolCallTransaction) sendSubagentMessage(ctx context.Context, input Sen if err != nil { return fmt.Errorf("marshal parent message metadata: %w", err) } - if _, err := createAgentContentInputTx(ctx, t.notifications, t.tx, t.q, CreateAgentContentInputInput{ + if _, err := createAgentContentInputTx(ctx, t.unit, t.q, CreateAgentContentInputInput{ ProjectID: child.ProjectID, AgentID: child.ID, Actor: actor, @@ -687,7 +559,7 @@ func LaunchSubagentForToolCall( if input.Subagent == nil || input.Subagent.ParentAgentID != tx.input.AgentID { return nil, errors.New("subagent launch must name the calling agent as parent") } - result, err := tx.store.launchAgentTx(ctx, tx.tx, tx.q, tx.notifications, input) + result, err := tx.store.launchAgentTx(ctx, tx.unit, tx.q, input) if err != nil { return nil, err } @@ -727,7 +599,14 @@ func StopSubagentForToolCall( } var machines []MachineRecord if input.Archive { - machines, err = archiveAgentTreeTx(ctx, tx.tx, tx.q, tx.notifications, child.ProjectID, child.ID, nil) + machines, err = archiveAgentTreeTx( + ctx, + tx.unit, + tx.q, + child.ProjectID, + child.ID, + nil, + ) if err != nil { return nil, err } @@ -748,10 +627,12 @@ func (t *toolCallTransaction) cancelSubagent(ctx context.Context, child AgentRec if child.State != AgentStateActive { return storeerr.InvalidRequest(errors.New("subagent is archived")) } - if err := lifecyclelock.Agents(ctx, t.tx, []lifecyclelock.AgentRef{{ - ProjectID: child.ProjectID, - AgentID: child.ID, - }}); err != nil { + if err := t.unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{ + {ProjectID: child.ProjectID, AgentID: child.ID}, + {ProjectID: t.input.ProjectID, AgentID: t.input.AgentID}, + }, + agentexecution.RuntimeAuthority{AgentID: t.input.AgentID, + RuntimeLockID: t.input.RuntimeLockID}); err != nil { return err } parent, err := loadAgentInProjectTx(ctx, t.tx, t.input.ProjectID, t.input.AgentID) @@ -766,13 +647,12 @@ func (t *toolCallTransaction) cancelSubagent(ctx context.Context, child AgentRec if err != nil { return err } - _, err = cancelAgentTx(ctx, t.notifications, t.tx, t.q, cancelAgentTxInput{ - ProjectID: child.ProjectID, - AgentID: child.ID, - ActorID: actorID, - ReasonCode: cancelReasonAgentCanceled, - ModelCallMessage: "The model call was canceled by the parent agent.", - }) + h, err := t.unit.Handle(child.ProjectID, child.ID) + if err != nil { + return err + } + _, err = h.Cancel(ctx, agentexecution.CancelInput{ActorID: actorID, Reason: cancelReasonAgentCanceled, + Message: "The model call was canceled by the parent agent."}) return err } @@ -785,7 +665,11 @@ func (s *Store) ArchiveIdleAgents(ctx context.Context, limit int) ([]MachineReco return s.archiveIdleAgents(ctx, nil, limit) } -func (s *Store) archiveIdleAgents(ctx context.Context, asOf *time.Time, limit int) ([]MachineRecord, int, error) { +func (s *Store) archiveIdleAgents( + ctx context.Context, + asOf *time.Time, + limit int, +) ([]MachineRecord, int, error) { if limit <= 0 { limit = 50 } @@ -812,9 +696,13 @@ func (s *Store) archiveIdleCandidates( var machines []MachineRecord archived := 0 for _, candidate := range candidates { - released, err := storeutil.RetryTransaction(ctx, "archive_idle_agent", func() ([]MachineRecord, error) { - return s.archiveIdleCandidateOnce(ctx, candidate, asOf) - }) + released, err := storeutil.RetryTransaction( + ctx, + "archive_idle_agent", + func() ([]MachineRecord, error) { + return s.archiveIdleCandidateOnce(ctx, candidate, asOf) + }, + ) if err != nil { if errors.Is(err, storeerr.ErrNotFound) || errors.Is(err, errIdleArchiveNoLongerEligible) { continue @@ -832,17 +720,18 @@ func (s *Store) archiveIdleCandidateOnce( candidate idleArchiveCandidate, asOf *time.Time, ) ([]MachineRecord, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return nil, fmt.Errorf("begin archive idle agent: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) if err := enterActiveAgentProjectTx(ctx, tx, qtx, candidate.ProjectID); err != nil { return nil, err } - locked, err := lockAgentTreeTx(ctx, tx, qtx, candidate.ProjectID, candidate.ID) + locked, err := lockAgentTreeTx(ctx, unit, qtx, candidate.ProjectID, candidate.ID) if err != nil { return nil, err } @@ -857,11 +746,11 @@ func (s *Store) archiveIdleCandidateOnce( if !idle { return nil, errIdleArchiveNoLongerEligible } - released, err := archiveLockedAgentTreeTx(ctx, tx, qtx, txNotifications, locked, nil) + released, err := archiveLockedAgentTreeTx(ctx, unit, qtx, locked, nil) if err != nil { return nil, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "archive idle agent"); err != nil { + if err := unit.Commit(ctx, "archive idle agent"); err != nil { return nil, err } return released, nil @@ -905,23 +794,39 @@ func (r *ToolCallReader) ReadSubagentTurnEvents( return SubagentStatus{}, nil, fmt.Errorf("check subagent turn ownership: %w", err) } if !owned { - return SubagentStatus{}, nil, storeerr.InvalidRequest(errors.New("no turn with that turn_id belongs to the subagent")) + return SubagentStatus{}, nil, storeerr.InvalidRequest( + errors.New("no turn with that turn_id belongs to the subagent"), + ) } - events, err := listTurnEventsForReadTx(ctx, r.transaction.q, projectID, status.AgentID, turnID, beforeSequence, limit) + events, err := listTurnEventsForReadTx( + ctx, + r.transaction.q, + projectID, + status.AgentID, + turnID, + beforeSequence, + limit, + ) if err != nil { return SubagentStatus{}, nil, err } return status, events, nil } -func (s *Store) ListSubagents(ctx context.Context, projectID, parentAgentID uuid.UUID) ([]SubagentStatus, error) { +func (s *Store) ListSubagents( + ctx context.Context, + projectID, parentAgentID uuid.UUID, +) ([]SubagentStatus, error) { if projectID == uuid.Nil || parentAgentID == uuid.Nil { return nil, errors.New("project and parent agent are required") } return listChildAgentsTx(ctx, s.q, projectID, parentAgentID, true) } -func (s *Store) ListAgentDescendantIDs(ctx context.Context, projectID, agentID uuid.UUID) ([]uuid.UUID, error) { +func (s *Store) ListAgentDescendantIDs( + ctx context.Context, + projectID, agentID uuid.UUID, +) ([]uuid.UUID, error) { if projectID == uuid.Nil || agentID == uuid.Nil { return nil, errors.New("project and agent are required") } diff --git a/internal/storage/executionstore/tool_call_commands.go b/internal/storage/executionstore/tool_call_commands.go index ccb9f9f33..35b76f2c6 100644 --- a/internal/storage/executionstore/tool_call_commands.go +++ b/internal/storage/executionstore/tool_call_commands.go @@ -25,7 +25,7 @@ func CompleteToolCallInExecution(input ToolCallCompletionInput) ToolCallCommand func StartToolCallAsync() ToolCallCommand { return toolCallCommandFunc(func(ctx context.Context, tx *toolCallTransaction) (any, error) { - return nil, tx.startToolCall(ctx, true, nil) + return nil, tx.startToolCall(ctx, true) }) } @@ -35,7 +35,7 @@ func StartProcessForToolCall(input CreateProcessInput) ToolCallCommand { if err != nil { return nil, err } - return record, tx.startToolCall(ctx, false, nil) + return record, tx.startToolCall(ctx, false) }) } @@ -45,7 +45,7 @@ func CreateProcessActionForToolCall(input CreateProcessActionInput) ToolCallComm if err != nil { return nil, err } - return record, tx.startToolCall(ctx, false, nil) + return record, tx.startToolCall(ctx, false) }) } @@ -68,7 +68,7 @@ func StopProcessForToolCall( if err != nil { return nil, err } - return record, tx.startToolCall(ctx, false, nil) + return record, tx.startToolCall(ctx, false) }) } @@ -78,7 +78,7 @@ func CreateQuestionForToolCall(input CreateQuestionInteractionInput) ToolCallCom if err != nil { return nil, err } - return record, tx.startToolCall(ctx, false, interactionUpdateFromRecord(record)) + return record, tx.startToolCall(ctx, false) }) } diff --git a/internal/storage/executionstore/tool_call_execution.go b/internal/storage/executionstore/tool_call_execution.go index 2b4dc77eb..3827802f0 100644 --- a/internal/storage/executionstore/tool_call_execution.go +++ b/internal/storage/executionstore/tool_call_execution.go @@ -8,7 +8,9 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -24,6 +26,7 @@ type ExecuteToolCallResult struct { Disposition ToolCallDisposition Applied bool CommandResult any + Completed *ToolCallRecord } type ToolCallDisposition uint8 @@ -45,8 +48,9 @@ type ToolCallReader struct { } type toolCallTransaction struct { + unit *agentexecution.Unit store *Store - tx pgx.Tx + tx dbsqlc.DBTX q *dbsqlc.Queries notifications *notifications.TxNotifications input ExecuteToolCallInput @@ -55,6 +59,8 @@ type toolCallTransaction struct { applied bool hasDurableCompletionOwner bool requiresWaitingDisposition bool + owner agentexecution.ToolOwner + completed *ToolCallRecord } func (s *Store) ExecuteToolCall( @@ -81,16 +87,18 @@ func (s *Store) executeToolCallOnce( input ExecuteToolCallInput, plan ToolCallPlan, ) (ExecuteToolCallResult, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return ExecuteToolCallResult{}, fmt.Errorf("begin tool call execution: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() toolTx := &toolCallTransaction{ + unit: unit, store: s, tx: tx, q: dbsqlc.New(tx), - notifications: s.newTxNotifications(), + notifications: unit.Notifications(), input: input, } command, err := plan(&ToolCallReader{transaction: toolTx}) @@ -119,13 +127,21 @@ func (s *Store) executeToolCallOnce( storeerr.ErrInvalidToolCallDisposition, ) } - if err := s.commitTxWithNotifications(ctx, tx, toolTx.notifications, "execute tool call"); err != nil { + if toolTx.disposition == ToolCallDispositionCompleted && toolTx.completed == nil { + record, err := getToolCallTx(ctx, tx, input.ProjectID, input.AgentID, input.ToolCallID) + if err != nil { + return ExecuteToolCallResult{}, err + } + toolTx.completed = &record + } + if err := unit.Commit(ctx, "execute tool call"); err != nil { return ExecuteToolCallResult{}, err } return ExecuteToolCallResult{ Disposition: toolTx.disposition, Applied: toolTx.applied, CommandResult: commandResult, + Completed: toolTx.completed, }, nil } @@ -137,107 +153,64 @@ func (t *toolCallTransaction) lockOrAcceptExisting(ctx context.Context) error { return t.lockToolCall(ctx, true) } -func (t *toolCallTransaction) lockToolCall( - ctx context.Context, - acceptExisting bool, -) error { +func (t *toolCallTransaction) lockToolCall(ctx context.Context, acceptExisting bool) error { if t == nil || t.tx == nil || t.q == nil { return errors.New("tool call transaction is required") } if t.locked || t.disposition != 0 { return nil } - if err := lockAgentRuntimeForOwnedMutationTx( - ctx, - t.q, - t.input.ProjectID, - t.input.AgentID, - t.input.RuntimeLockID, - ); err != nil { + err := t.unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: t.input.ProjectID, AgentID: t.input.AgentID}}, + agentexecution.IngressAuthority{}, + ) + if errors.Is(err, storeerr.ErrNotFound) { + return storeerr.ErrRuntimeLockInactive + } + if err != nil { return err } - existing, err := t.q.GetToolCallDispatchState( - ctx, - dbsqlc.GetToolCallDispatchStateParams{ - ProjectID: t.input.ProjectID, - AgentID: t.input.AgentID, - ID: t.input.ToolCallID, - }, + h, err := t.unit.Handle(t.input.ProjectID, t.input.AgentID) + if err != nil { + return err + } + dispatch, err := h.DispatchTool(ctx, + agentexecution.ToolRef{ID: t.input.ToolCallID, RuntimeLockID: t.input.RuntimeLockID}, acceptExisting, ) if err != nil { - return fmt.Errorf("load tool call before execution: %w", err) + return err } - switch ToolCallState(existing.State) { - case ToolCallStateReady: + switch dispatch { + case agentexecution.ToolDispatchReady: t.locked = true - return nil - case ToolCallStateAwaitingAuthorization, ToolCallStateAwaitingPermission: - return storeerr.ErrIdempotencyConflict - case ToolCallStateRunning: - return storeerr.ErrToolCallInProgress - case ToolCallStateWaiting: - if acceptExisting { - t.disposition = ToolCallDispositionWaiting - return nil - } - case ToolCallStateCompleted: - if acceptExisting { - t.disposition = ToolCallDispositionCompleted - return nil - } + case agentexecution.ToolDispatchWaiting: + t.disposition = ToolCallDispositionWaiting + case agentexecution.ToolDispatchCompleted: + t.disposition = ToolCallDispositionCompleted } - return storeerr.ErrIdempotencyConflict + return nil } -func (t *toolCallTransaction) startToolCall( - ctx context.Context, - retainRuntimeOwnership bool, - interactionUpdate *notifications.AgentInteractionUpdate, -) error { - if t == nil || t.tx == nil || t.q == nil { - return errors.New("tool call transaction is required") - } +func (t *toolCallTransaction) startToolCall(ctx context.Context, retainRuntimeOwnership bool) error { if t.disposition != 0 { return nil } if err := t.lockForMutation(ctx); err != nil { return err } - changed, err := t.q.StartToolCall( - ctx, - dbsqlc.StartToolCallParams{ - RetainRuntimeOwnership: retainRuntimeOwnership, - ProjectID: t.input.ProjectID, - AgentID: t.input.AgentID, - ID: t.input.ToolCallID, - RuntimeLockID: t.input.RuntimeLockID, - }, - ) + h, err := t.unit.Handle(t.input.ProjectID, t.input.AgentID) if err != nil { - return fmt.Errorf("start tool call: %w", err) - } - if changed == 0 { - if runtimeErr := agentRuntimeLockActiveTx( - ctx, - t.q, - t.input.ProjectID, - t.input.AgentID, - t.input.RuntimeLockID, - ); runtimeErr != nil { - return runtimeErr - } - return storeerr.ErrIdempotencyConflict + return err } - state := ToolCallStateWaiting + ref := agentexecution.ToolRef{ID: t.input.ToolCallID, RuntimeLockID: t.input.RuntimeLockID} if retainRuntimeOwnership { + t.applied, err = h.RunTool(ctx, ref) t.disposition = ToolCallDispositionRunning - state = ToolCallStateRunning } else { + t.applied, err = h.WaitForTool(ctx, ref, t.owner) t.disposition = ToolCallDispositionWaiting } - t.notifications.AddToolCallUpdate(t.input.AgentID, t.input.ToolCallID, string(state), interactionUpdate) - t.applied = true - return nil + return err } func (t *toolCallTransaction) completeToolCall( @@ -250,7 +223,7 @@ func (t *toolCallTransaction) completeToolCall( if t.disposition != 0 { return ToolCallRecord{}, storeerr.ErrStateTransitionConflict } - record, err := completeToolCallTx(ctx, t.notifications, t.tx, CompleteToolCallInput{ + record, err := completeExecutionTool(ctx, t.unit, CompleteToolCallInput{ ProjectID: t.input.ProjectID, AgentID: t.input.AgentID, ID: t.input.ToolCallID, @@ -264,6 +237,7 @@ func (t *toolCallTransaction) completeToolCall( t.locked = true t.disposition = ToolCallDispositionCompleted t.applied = true + t.completed = &record return record, nil } diff --git a/internal/storage/executionstore/tool_call_reads_store.go b/internal/storage/executionstore/tool_call_reads_store.go index 63f67c1dc..2d2861d57 100644 --- a/internal/storage/executionstore/tool_call_reads_store.go +++ b/internal/storage/executionstore/tool_call_reads_store.go @@ -14,7 +14,7 @@ import ( func getToolCallTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID, id uuid.UUID, ) (ToolCallRecord, error) { if projectID == uuid.Nil || agentID == uuid.Nil || id == uuid.Nil { diff --git a/internal/storage/executionstore/tool_calls_store.go b/internal/storage/executionstore/tool_calls_store.go index 722d4bd29..81d2d07c5 100644 --- a/internal/storage/executionstore/tool_calls_store.go +++ b/internal/storage/executionstore/tool_calls_store.go @@ -8,12 +8,11 @@ import ( "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/events" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/listing" "github.com/omnara-ai/omnara/internal/storage/storeerr" - "github.com/omnara-ai/omnara/internal/toolcatalog" ) type CompleteToolCallInput struct { @@ -148,14 +147,14 @@ func (s *Store) CompleteToolCall( } input.ResultContentParts = parts record, err := s.completeToolCallOnce(ctx, input) - var read *toolResultArtifactReadRequiredError - if !errors.As(err, &read) { - return record, err + if err == nil { + return record, nil } - ctx, err = s.loadToolResultForReplay(ctx, read) + parts, err = s.replayToolContent(ctx, input.ProjectID, input.AgentID, input.ResultContentParts, err) if err != nil { return ToolCallRecord{}, err } + input.ResultContentParts = parts return s.completeToolCallOnce(ctx, input) } @@ -163,277 +162,94 @@ func (s *Store) completeToolCallOnce( ctx context.Context, input CompleteToolCallInput, ) (ToolCallRecord, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return ToolCallRecord{}, fmt.Errorf("begin complete tool call: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() - record, err := completeToolCallTx(ctx, txNotifications, tx, input) + defer func() { _ = unit.Rollback(ctx) }() + + record, err := completeExecutionTool(ctx, unit, input) if err != nil { return ToolCallRecord{}, err } - if err := s.commitTxWithNotifications(ctx, tx, txNotifications, "complete tool call"); err != nil { + if err := unit.Commit(ctx, "complete tool call"); err != nil { return ToolCallRecord{}, err } return record, nil } -func (s *Store) MarkToolCallReady( - ctx context.Context, - input MarkToolCallReadyInput, -) (ToolCallRecord, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ID == uuid.Nil || - input.RuntimeLockID == uuid.Nil { - return ToolCallRecord{}, errors.New( - "project, agent, tool call id, and runtime lock are required", - ) - } - txNotifications := s.newTxNotifications() - tx, qtx, err := s.beginAgentRuntimeOwnedMutation( - ctx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ) +func (s *Store) MarkToolCallReady(ctx context.Context, input MarkToolCallReadyInput) (ToolCallRecord, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { return ToolCallRecord{}, err } - defer func() { _ = tx.Rollback(ctx) }() - row, err := qtx.MarkToolCallReady( - ctx, - dbsqlc.MarkToolCallReadyParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - RuntimeLockID: input.RuntimeLockID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - existing, loadErr := getToolCallTx(ctx, tx, input.ProjectID, input.AgentID, input.ID) - if loadErr != nil { - return ToolCallRecord{}, loadErr - } - if err := agentRuntimeLockActiveTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { - return ToolCallRecord{}, err - } - if existing.State == ToolCallStateReady && existing.CompletedAt == nil { - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "mark tool call ready", - ); err != nil { - return ToolCallRecord{}, err - } - return existing, nil - } - return ToolCallRecord{}, storeerr.ErrStateTransitionConflict + defer func() { _ = unit.Rollback(ctx) }() + if _, + err = h.AuthorizeTool(ctx, + agentexecution.ToolRef{ID: input.ID, + RuntimeLockID: input.RuntimeLockID}); err != nil { + return ToolCallRecord{}, err } + record, err := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, input.ID) if err != nil { - return ToolCallRecord{}, fmt.Errorf("mark tool call ready: %w", err) + return ToolCallRecord{}, err } - record := toolCallRecordFromReadySQLC(row) - txNotifications.AddToolCallUpdate(record.AgentID, record.ID, string(record.State), nil) - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "mark tool call ready", - ); err != nil { + if err = unit.Commit(ctx, "authorize tool"); err != nil { return ToolCallRecord{}, err } return record, nil } -func (s *Store) RequeueRuntimeToolCall( - ctx context.Context, - input RequeueRuntimeToolCallInput, -) error { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ToolCallID == uuid.Nil || - input.RuntimeLockID == uuid.Nil { - return errors.New("project, agent, tool call, and runtime lock are required") - } - txNotifications := s.newTxNotifications() - tx, qtx, err := s.beginAgentRuntimeOwnedMutation( - ctx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ) +func (s *Store) RequeueRuntimeToolCall(ctx context.Context, input RequeueRuntimeToolCallInput) error { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) if err != nil { return err } - defer func() { _ = tx.Rollback(ctx) }() - changed, err := qtx.RequeueRuntimeToolCall( - ctx, - dbsqlc.RequeueRuntimeToolCallParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ToolCallID, - RuntimeLockID: input.RuntimeLockID, - }, - ) - if err != nil { - return fmt.Errorf("requeue runtime tool call: %w", err) - } - if changed == 0 { - existing, err := qtx.GetToolCallDispatchState( - ctx, - dbsqlc.GetToolCallDispatchStateParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ToolCallID, - }, - ) - if err != nil { - return fmt.Errorf("load tool call after requeue miss: %w", err) - } - if existing.State != string(ToolCallStateReady) || - existing.RuntimeLockID != nil { - return storeerr.ErrStateTransitionConflict - } - } else { - txNotifications.AddToolCallUpdate( - input.AgentID, - input.ToolCallID, - string(ToolCallStateReady), - nil, - ) - } - metadata, err := marshalJSON(map[string]any{ - "reason": "async_preflight_retry", - "tool_call_id": input.ToolCallID, - }) - if err != nil { - return fmt.Errorf("marshal requeued runtime tool wakeup: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - Metadata: metadata, - }, - ); err != nil { - return fmt.Errorf("mark requeued runtime tool wakeup: %w", err) - } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "requeue runtime tool call", - ); err != nil { + defer func() { _ = unit.Rollback(ctx) }() + if _, + err = h.RequeueTool(ctx, + agentexecution.ToolRef{ID: input.ToolCallID, + RuntimeLockID: input.RuntimeLockID}); err != nil { return err } - return nil + return unit.Commit(ctx, "requeue tool") } -func completeToolCallTx( +func completeExecutionTool( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, input CompleteToolCallInput, ) (ToolCallRecord, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ID == uuid.Nil || - input.RuntimeLockID == uuid.Nil { - return ToolCallRecord{}, errors.New( - "project, agent, tool call id, and runtime lock are required", - ) - } - if !input.Outcome.IsTerminal() { - return ToolCallRecord{}, fmt.Errorf("invalid tool result outcome %q", input.Outcome) + + err := unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, AgentID: input.AgentID}}, + agentexecution.IngressAuthority{}, + ) + if errors.Is(err, storeerr.ErrNotFound) { + return ToolCallRecord{}, storeerr.ErrRuntimeLockInactive } - if len(input.ResultContentParts) == 0 { - input.ResultContentParts = json.RawMessage(`[]`) + if err != nil { + return ToolCallRecord{}, err } + blocks, err := parseToolResultContentBlocks(input.ResultContentParts) if err != nil { - return ToolCallRecord{}, fmt.Errorf("invalid model-visible content parts: %w", err) + return ToolCallRecord{}, err } - normalizedParts, err := marshalToolResultContentBlocks(blocks) + parts, err := executionContent(blocks) if err != nil { - return ToolCallRecord{}, fmt.Errorf("canonicalize model-visible content parts: %w", err) - } - input.ResultContentParts = normalizedParts - qtx := dbsqlc.New(tx) - if err := lockAgentRuntimeForOwnedMutationTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return ToolCallRecord{}, err } - row, err := qtx.CompleteToolCall( - ctx, - dbsqlc.CompleteToolCallParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - RuntimeLockID: input.RuntimeLockID, - Outcome: string(input.Outcome), - }) - if err == nil { - record := toolCallRecordFromCompleteSQLC(row) - record.ResultContentParts = input.ResultContentParts - admitted, err := appendToolResultRecordTx( - ctx, - txNotifications, - tx, - record, - nil, - ) - if err != nil { - return ToolCallRecord{}, err - } - applyAdmittedToolResult(&record, admitted) - return record, nil - } - if !errors.Is(err, pgx.ErrNoRows) { - return ToolCallRecord{}, fmt.Errorf("complete tool call: %w", err) - } - existing, loadErr := getToolCallTx(ctx, tx, input.ProjectID, input.AgentID, input.ID) - if loadErr != nil { - return ToolCallRecord{}, loadErr + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return ToolCallRecord{}, err } - if err := agentRuntimeLockActiveTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { + _, err = h.CompleteTool(ctx, agentexecution.ToolCompletion{ToolRef: agentexecution.ToolRef{ID: input.ID, + RuntimeLockID: input.RuntimeLockID}, Outcome: string(input.Outcome), Content: parts}) + if err != nil { return ToolCallRecord{}, err } - if toolcatalog.IsPlatformManagedToolType(existing.Type) && - existing.State == ToolCallStateCompleted && - existing.Outcome == input.Outcome { - ok, err := completedToolCallMatchesTx( - ctx, - dbsqlc.New(tx), - input.ProjectID, - input.AgentID, - input.ID, - input.Outcome, - input.ResultContentParts, - ) - if err != nil { - return ToolCallRecord{}, err - } - if !ok { - return ToolCallRecord{}, storeerr.ErrIdempotencyConflict - } - return existing, nil - } - return ToolCallRecord{}, storeerr.ErrStateTransitionConflict + return getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, input.ID) } func (s *Store) CompleteRuntimeToolCall( @@ -446,14 +262,14 @@ func (s *Store) CompleteRuntimeToolCall( } input.ResultContentParts = parts record, err := s.completeRuntimeToolCallOnce(ctx, input) - var read *toolResultArtifactReadRequiredError - if !errors.As(err, &read) { - return record, err + if err == nil { + return record, nil } - ctx, err = s.loadToolResultForReplay(ctx, read) + parts, err = s.replayToolContent(ctx, input.ProjectID, input.AgentID, input.ResultContentParts, err) if err != nil { return ToolCallRecord{}, err } + input.ResultContentParts = parts return s.completeRuntimeToolCallOnce(ctx, input) } @@ -461,173 +277,28 @@ func (s *Store) completeRuntimeToolCallOnce( ctx context.Context, input CompleteRuntimeToolCallInput, ) (ToolCallRecord, error) { - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) - if err != nil { - return ToolCallRecord{}, err - } - defer func() { _ = tx.Rollback(ctx) }() - record, err := completeRuntimeToolCallTx(ctx, txNotifications, tx, input) - if err != nil { - return ToolCallRecord{}, err - } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "complete runtime tool call", - ); err != nil { - return ToolCallRecord{}, err - } - return record, nil -} - -func completeRuntimeToolCallTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - input CompleteRuntimeToolCallInput, -) (ToolCallRecord, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ID == uuid.Nil || - input.RuntimeLockID == uuid.Nil { - return ToolCallRecord{}, errors.New( - "project, agent, tool call id, and runtime lock are required", - ) - } if !input.Outcome.IsTerminal() || input.Outcome == ToolResultOutcomeDenied { - return ToolCallRecord{}, fmt.Errorf( - "invalid runtime tool result outcome %q", - input.Outcome, - ) + return ToolCallRecord{}, errors.New("invalid runtime tool outcome") } - if len(input.ResultContentParts) == 0 { - input.ResultContentParts = json.RawMessage(`[]`) - } - blocks, err := parseToolResultContentBlocks(input.ResultContentParts) + unit, err := s.cell.Begin(ctx) if err != nil { - return ToolCallRecord{}, fmt.Errorf("invalid model-visible content parts: %w", err) - } - normalizedParts, err := marshalToolResultContentBlocks(blocks) - if err != nil { - return ToolCallRecord{}, fmt.Errorf("canonicalize model-visible content parts: %w", err) - } - input.ResultContentParts = normalizedParts - qtx := dbsqlc.New(tx) - if err := lockAgentRuntimeForOwnedMutationTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { return ToolCallRecord{}, err } - row, err := qtx.CompleteRuntimeToolCall( + defer func() { _ = unit.Rollback(ctx) }() + record, err := completeExecutionTool( ctx, - dbsqlc.CompleteRuntimeToolCallParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - RuntimeLockID: input.RuntimeLockID, - Outcome: string(input.Outcome), - }, - ) - if err == nil { - return finishCompletedToolCallTx( - ctx, - txNotifications, - tx, - qtx, - toolCallRecordFromRuntimeCompleteSQLC(row), - toolCallResultInput{ - Outcome: input.Outcome, - ResultContentParts: input.ResultContentParts, - }, - ) - } - if !errors.Is(err, pgx.ErrNoRows) { - return ToolCallRecord{}, fmt.Errorf("complete runtime tool call: %w", err) - } - existing, loadErr := getToolCallTx(ctx, tx, input.ProjectID, input.AgentID, input.ID) - if loadErr != nil { - return ToolCallRecord{}, loadErr - } - if err := agentRuntimeLockActiveTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.RuntimeLockID, - ); err != nil { - return ToolCallRecord{}, err - } - if toolcatalog.IsPlatformManagedToolType(existing.Type) && - existing.State == ToolCallStateCompleted && - existing.Outcome == input.Outcome { - ok, err := completedToolCallMatchesTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.ID, - input.Outcome, - input.ResultContentParts, - ) - if err != nil { - return ToolCallRecord{}, err - } - if ok { - return existing, nil - } - } - return ToolCallRecord{}, storeerr.ErrIdempotencyConflict -} - -type toolCallResultInput struct { - Outcome ToolResultOutcome - ResultContentParts json.RawMessage -} - -func finishCompletedToolCallTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - qtx *dbsqlc.Queries, - record ToolCallRecord, - input toolCallResultInput, -) (ToolCallRecord, error) { - record.ResultContentParts = input.ResultContentParts - admitted, err := appendToolResultRecordTx( - ctx, - txNotifications, - tx, - record, - nil, + unit, + CompleteToolCallInput{ProjectID: input.ProjectID, AgentID: input.AgentID, + ID: input.ID, + RuntimeLockID: input.RuntimeLockID, + Outcome: input.Outcome, + ResultContentParts: input.ResultContentParts}, ) if err != nil { return ToolCallRecord{}, err } - applyAdmittedToolResult(&record, admitted) - metadata, err := marshalJSON( - map[string]any{ - "reason": "tool_result", - "tool_call_id": record.ID, - "model_call_context_id": record.ModelCallContextID, - "outcome": input.Outcome, - }, - ) - if err != nil { - return ToolCallRecord{}, fmt.Errorf("marshal runtime tool result wakeup metadata: %w", err) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - Metadata: metadata, - }, - ); err != nil { - return ToolCallRecord{}, fmt.Errorf("mark runtime tool result wakeup: %w", err) + if err = unit.Commit(ctx, "complete runtime tool"); err != nil { + return ToolCallRecord{}, err } return record, nil } @@ -641,22 +312,17 @@ func (s *Store) CompleteCustomToolCall( return CompleteCustomToolCallResult{}, err } input.ContentBlocks = parts - txNotifications := s.newTxNotifications() - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return CompleteCustomToolCallResult{}, err } - defer func() { _ = tx.Rollback(ctx) }() - result, err := completeCustomToolCallTx(ctx, txNotifications, tx, input) + defer func() { _ = unit.Rollback(ctx) }() + + result, err := completeCustomToolCallTx(ctx, unit, input) if err != nil { return CompleteCustomToolCallResult{}, err } - if err := s.commitTxWithNotifications( - ctx, - tx, - txNotifications, - "complete custom tool call", - ); err != nil { + if err := unit.Commit(ctx, "complete custom tool call"); err != nil { return CompleteCustomToolCallResult{}, err } return result, nil @@ -664,101 +330,40 @@ func (s *Store) CompleteCustomToolCall( func completeCustomToolCallTx( ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, + unit *agentexecution.Unit, input CompleteCustomToolCallInput, ) (CompleteCustomToolCallResult, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ID == uuid.Nil { - return CompleteCustomToolCallResult{}, errors.New( - "project, agent, and tool call ids are required", - ) - } - switch input.Outcome { - case ToolResultOutcomeSucceeded, ToolResultOutcomeFailed: - default: - return CompleteCustomToolCallResult{}, errors.New( - "custom tool result outcome must be succeeded or failed", - ) + if err := unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, + AgentID: input.AgentID}}, + agentexecution.ExternalAuthority{}); err != nil { + return CompleteCustomToolCallResult{}, err } blocks, err := parseToolResultContentBlocks(input.ContentBlocks) if err != nil { return CompleteCustomToolCallResult{}, err } - resultContentParts, err := marshalToolResultContentBlocks(blocks) + parts, err := executionContent(blocks) if err != nil { return CompleteCustomToolCallResult{}, err } - qtx := dbsqlc.New(tx) - row, err := qtx.CompleteCustomToolCall( - ctx, - dbsqlc.CompleteCustomToolCallParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - Outcome: string(input.Outcome), - }, - ) - if err == nil { - record := toolCallRecordFromCustomCompleteSQLC(row) - record.ResultContentParts = resultContentParts - admitted, err := appendToolResultRecordTx( - ctx, - txNotifications, - tx, - record, - nil, - ) - if err != nil { - return CompleteCustomToolCallResult{}, err - } - if !admitted.Inserted { - return CompleteCustomToolCallResult{}, storeerr.ErrIdempotencyConflict - } - applyAdmittedToolResult(&record, admitted) - metadata, err := marshalJSON( - map[string]any{ - "reason": "custom_tool_result", - "tool_call_id": record.ID, - "model_call_context_id": record.ModelCallContextID, - }, - ) - if err != nil { - return CompleteCustomToolCallResult{}, fmt.Errorf( - "marshal custom tool result wakeup metadata: %w", - err, - ) - } - if err := qtx.MarkAgentWakeup( - ctx, - dbsqlc.MarkAgentWakeupParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - Metadata: metadata, - }, - ); err != nil { - return CompleteCustomToolCallResult{}, fmt.Errorf( - "mark custom tool result wakeup: %w", - err, - ) - } - return CompleteCustomToolCallResult{ - ToolCall: record, - Event: admitted.Event, - ContentBlocks: admitted.ContentBlocks, - }, nil - } - if !errors.Is(err, pgx.ErrNoRows) { - return CompleteCustomToolCallResult{}, fmt.Errorf( - "complete custom tool call: %w", - err, - ) - } - existing, loadErr := getToolCallTx(ctx, tx, input.ProjectID, input.AgentID, input.ID) - if loadErr != nil { - return CompleteCustomToolCallResult{}, loadErr - } - if existing.Type != toolcatalog.ToolTypeCustom { - return CompleteCustomToolCallResult{}, storeerr.ErrNotFound - } - return CompleteCustomToolCallResult{}, storeerr.ErrStateTransitionConflict + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return CompleteCustomToolCallResult{}, err + } + completed, err := h.CompleteCustomTool(ctx, input.ID, string(input.Outcome), parts) + if err != nil { + return CompleteCustomToolCallResult{}, err + } + record, err := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, input.ID) + if err != nil { + return CompleteCustomToolCallResult{}, err + } + return CompleteCustomToolCallResult{ToolCall: record, Event: TypedAgentEventRecord{ + Event: executionEvent( + input.AgentID, + completed.Event, + events.KindToolResult, + ), TurnID: completed.Event.TurnID, + ToolCallResultID: completed.ResultID}, ContentBlocks: record.ResultContentParts}, nil } diff --git a/internal/storage/executionstore/tool_calls_test_helpers_test.go b/internal/storage/executionstore/tool_calls_test_helpers_test.go index f72c2e100..aa0160969 100644 --- a/internal/storage/executionstore/tool_calls_test_helpers_test.go +++ b/internal/storage/executionstore/tool_calls_test_helpers_test.go @@ -43,7 +43,7 @@ func createReadyToolCallsForTest( if err != nil { t.Fatalf("create tool-call seed input: %v", err) } - admitted, found := admitNextAgentInputAndOpenTurnForTest( + _, found := admitNextAgentInputAndOpenTurnForTest( t, ctx, store, @@ -54,17 +54,17 @@ func createReadyToolCallsForTest( if !found { t.Fatal("expected tool-call seed input admission") } - claim, err := store.Execution().ClaimNormalModelCall( + prepared1, err := store.Execution().PrepareNormalModelCall( ctx, - executionstore.ClaimNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: agentID, - RuntimeLockID: lock.ID, - OpeningInputIDs: []uuid.UUID{input.ID}, - AgentConfigID: configID, - InputEventSequence: admitted.Events[0].Sequence, + executionstore.PrepareNormalModelCallInput{ + ProjectID: testProjectID, + AgentID: agentID, + RuntimeLockID: lock.ID, + OpeningInputIDs: []uuid.UUID{input.ID}, }, ) + claim := prepared1.Claim + if err != nil { t.Fatalf("claim tool-call seed model context: %v", err) } diff --git a/internal/storage/executionstore/tool_calls_workflow.go b/internal/storage/executionstore/tool_calls_workflow.go index 4ef9419c9..d65b456e6 100644 --- a/internal/storage/executionstore/tool_calls_workflow.go +++ b/internal/storage/executionstore/tool_calls_workflow.go @@ -10,226 +10,12 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/daemonprotocol" - "github.com/omnara-ai/omnara/internal/events" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/processcmd" "github.com/omnara-ai/omnara/internal/publicid" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/toolcatalog" ) -const ( - toolCallResultEventIDPrefix = "tool-call-result-authority:" - toolCallCompletedInteractionReason = "tool_call_completed" -) - -func appendToolResultEventTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - record ToolCallRecord, - interactionUpdate *notifications.AgentInteractionUpdate, -) (events.Event, error) { - admitted, err := appendToolResultRecordTx( - ctx, - txNotifications, - tx, - record, - interactionUpdate, - ) - if err != nil { - return events.Event{}, err - } - return admitted.Event.Event, nil -} - -func appendToolResultRecordTx( - ctx context.Context, - txNotifications *notifications.TxNotifications, - tx pgx.Tx, - record ToolCallRecord, - interactionUpdate *notifications.AgentInteractionUpdate, -) (admittedToolCallResult, error) { - interactions, err := dbsqlc.New(tx).CancelOpenAgentInteractionsForToolCall( - ctx, - dbsqlc.CancelOpenAgentInteractionsForToolCallParams{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - ToolCallID: record.ID, - Reason: toolCallCompletedInteractionReason, - }, - ) - if err != nil { - return admittedToolCallResult{}, fmt.Errorf( - "close interactions for completed tool call: %w", - err, - ) - } - for _, interaction := range interactions { - interactionUpdate = ¬ifications.AgentInteractionUpdate{ - ID: interaction.ID, InteractionKind: interaction.InteractionKind, State: interaction.State, - } - } - admitted, err := admitToolCallResultTx( - ctx, - txNotifications, - tx, - CreateToolCallResultAuthorityInput{ - ProjectID: record.ProjectID, - AgentID: record.AgentID, - TurnID: record.TurnID, - ToolCallID: record.ID, - Outcome: record.Outcome, - ResultContentParts: record.ResultContentParts, - IdempotencyKey: toolCallResultEventIDPrefix + record.ID.String(), - }, - ) - if err != nil { - return admittedToolCallResult{}, err - } - if admitted.Event.ToolCallResultID != admitted.Result.ID { - return admittedToolCallResult{}, storeerr.ErrIdempotencyConflict - } - if admitted.Inserted { - if err := updateAgentTurnLatestEventTx( - ctx, - tx, - record.ProjectID, - record.AgentID, - record.TurnID, - admitted.Event.Event.ID, - uuid.Nil, - ); err != nil { - return admittedToolCallResult{}, err - } - txNotifications.AddToolCallUpdate( - record.AgentID, - record.ID, - string(ToolCallStateCompleted), - interactionUpdate, - ) - } - return admitted, nil -} - -func applyAdmittedToolResult(record *ToolCallRecord, admitted admittedToolCallResult) { - record.ResultContentParts = admitted.ContentBlocks - record.CompletedAt = timePtr(admitted.Result.CompletedAt) -} - -func completedToolCallMatchesTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, toolCallID uuid.UUID, - outcome ToolResultOutcome, - contentParts json.RawMessage, -) (bool, error) { - blocks, err := parseToolResultContentBlocks(contentParts) - if err != nil { - return false, err - } - contentParts, err = marshalToolResultContentBlocks(blocks) - if err != nil { - return false, err - } - existing, err := qtx.GetToolCall( - ctx, - dbsqlc.GetToolCallParams{ProjectID: projectID, AgentID: agentID, ID: toolCallID}, - ) - if err != nil { - return false, fmt.Errorf("load linked tool call for replay validation: %w", err) - } - if existing.State != "completed" || existing.CompletedAt == nil { - return false, nil - } - result, err := qtx.GetToolCallResultByToolCall( - ctx, - dbsqlc.GetToolCallResultByToolCallParams{ - ProjectID: projectID, - AgentID: agentID, - ToolCallID: toolCallID, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return false, nil - } - if err != nil { - return false, fmt.Errorf("load linked tool result for replay validation: %w", err) - } - storedParts, err := toolCallResultContentBlocksTx( - ctx, - qtx, - projectID, - agentID, - result.ID, - ) - if err != nil { - return false, err - } - if result.Outcome != string(outcome) { - return false, nil - } - return toolResultContentMatchesTx(ctx, qtx, projectID, agentID, toolCallID, storedParts, blocks, contentParts) -} - -func processToolResult(process ProcessRecord) (ToolResultOutcome, json.RawMessage, error) { - if !isProcessTerminal(process.State) { - return "", nil, fmt.Errorf("process %s is not terminal", process.ID) - } - exitCode := any(nil) - if process.ExitCode != nil { - exitCode = *process.ExitCode - } - errText := process.StateReasonMessage - if errText == "" { - errText = process.StateReasonCode - } - succeeded := process.State == ProcessStateExited && process.ExitCode != nil && *process.ExitCode == 0 - if process.State == ProcessStateExited && !succeeded && errText == "" && process.ExitCode != nil { - errText = fmt.Sprintf("exit code %d", *process.ExitCode) - } - result, err := marshalJSON(map[string]any{ - "process_id": publicResourceID(publicid.KindProcess, process.ID), - "exit_code": exitCode, - "state": process.State, - "state_reason_code": process.StateReasonCode, - "error": errText, - }) - if err != nil { - return "", nil, fmt.Errorf("marshal process tool result: %w", err) - } - outcome := ToolResultOutcomeSucceeded - if !succeeded { - outcome = ToolResultOutcomeFailed - } - return outcome, result, nil -} - -func startedProcessToolResult(process ProcessRecord, observed json.RawMessage) (json.RawMessage, error) { - commandLabel := process.ExecutionSpec.Label() - result := map[string]any{ - "process_id": publicResourceID(publicid.KindProcess, process.ID), - "state": process.State, - "command": commandLabel, - "next_action": "use write_process for stdin or read_process for retained output", - } - if len(observed) != 0 && string(observed) != "null" && string(observed) != "{}" { - var observedObject map[string]any - if err := json.Unmarshal(observed, &observedObject); err != nil { - return nil, fmt.Errorf("decode process started result: %w", err) - } - for key, value := range observedObject { - result[key] = value - } - result["process_id"] = publicResourceID(publicid.KindProcess, process.ID) - result["state"] = process.State - result["command"] = commandLabel - result["next_action"] = "use write_process for stdin or read_process for retained output" - } - return marshalJSON(result) -} - func commandTerminalToolResult( processID uuid.UUID, result json.RawMessage, @@ -347,9 +133,3 @@ func ToolResultContentParts(result json.RawMessage) (json.RawMessage, error) { } return marshalJSON([]map[string]any{{"type": "structured_data", "value": result}}) } - -func canceledToolResultContentParts() (json.RawMessage, error) { - return ToolResultContentParts( - json.RawMessage(`{"reason":"Agent canceled before this tool call completed."}`), - ) -} diff --git a/internal/storage/executionstore/tool_calls_workflow_test.go b/internal/storage/executionstore/tool_calls_workflow_test.go index 6e974cf1c..81720a207 100644 --- a/internal/storage/executionstore/tool_calls_workflow_test.go +++ b/internal/storage/executionstore/tool_calls_workflow_test.go @@ -9,108 +9,13 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/processcmd" + "github.com/omnara-ai/omnara/internal/processresult" "github.com/omnara-ai/omnara/internal/publicid" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/toolcatalog" "github.com/stretchr/testify/require" ) -func TestBindToolCallsUsesProviderEnvelopeAsCanonicalSource(t *testing.T) { - id := uuid.New() - envelope := modelenvelope.ResponseEnvelope{ - Normalized: modelenvelope.ResponseNormalized{ - Content: []modelenvelope.ResponsePart{ - {Type: modelenvelope.ResponsePartTypeText, Text: "before"}, - { - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_1", - ToolName: "lookup_customer", - ToolInput: json.RawMessage(`{"email":"ada@example.com"}`), - }, - }, - }, - } - calls, err := bindToolCalls(envelope, []ToolCallBindingInput{{ - ID: id, - ProviderCallID: "call_1", - Type: toolcatalog.ToolTypeCustom, - }}) - if err != nil { - t.Fatalf("bind tool calls: %v", err) - } - if len(calls) != 1 || - calls[0].ID != id || - calls[0].ProviderCallID != "call_1" || - calls[0].Name != "lookup_customer" || - string(calls[0].Input) != `{"email":"ada@example.com"}` || - calls[0].Type != toolcatalog.ToolTypeCustom { - t.Fatalf("bound tool call = %+v", calls) - } -} - -func TestBindToolCallsRequiresExactProviderCallCoverage(t *testing.T) { - envelope := modelenvelope.ResponseEnvelope{ - Normalized: modelenvelope.ResponseNormalized{ - Content: []modelenvelope.ResponsePart{{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_1", - ToolName: "run_command", - ToolInput: json.RawMessage(`{"command":"true"}`), - }}, - }, - } - for _, bindings := range [][]ToolCallBindingInput{ - {{ - ProviderCallID: "call_other", - Type: toolcatalog.ToolTypeBuiltIn, - }}, - { - { - ProviderCallID: "call_1", - Type: toolcatalog.ToolTypeBuiltIn, - }, - { - ProviderCallID: "call_extra", - Type: toolcatalog.ToolTypeBuiltIn, - }, - }, - } { - if _, err := bindToolCalls(envelope, bindings); err == nil { - t.Fatalf("bindings %+v did not fail", bindings) - } - } -} - -func TestBindToolCallsRejectsInvalidToolInput(t *testing.T) { - for _, input := range []json.RawMessage{ - json.RawMessage(`{"command":`), - json.RawMessage(`null`), - json.RawMessage(`[]`), - json.RawMessage(`"command"`), - } { - envelope := modelenvelope.ResponseEnvelope{ - Normalized: modelenvelope.ResponseNormalized{ - Content: []modelenvelope.ResponsePart{{ - Type: modelenvelope.ResponsePartTypeToolCall, - ProviderCallID: "call_1", - ToolName: "run_command", - ToolInput: input, - }}, - }, - } - _, err := bindToolCalls(envelope, []ToolCallBindingInput{{ - ID: uuid.New(), - ProviderCallID: "call_1", - Type: toolcatalog.ToolTypeBuiltIn, - }}) - if err == nil { - t.Fatalf("bindToolCalls accepted input %s", input) - } - } -} - func TestStructuredToolResultKeepsEmptyObject(t *testing.T) { parts, err := ToolResultContentParts(json.RawMessage(`{}`)) if err != nil { @@ -169,8 +74,8 @@ func TestStartedProcessToolResultKeepsProcessFactsAuthoritative(t *testing.T) { ID: uuid.New(), State: ProcessStateRunning, } - result, err := startedProcessToolResult( - process, + result, err := processresult.Started( + process.resultFacts(), json.RawMessage(`{"state":"exited","command":"other","next_action":"stop","output":"ready"}`), ) if err != nil { @@ -214,7 +119,12 @@ func TestFileTransferToolResultContentParts(t *testing.T) { {name: "wrong path", metadata: `{"path":"/memory/team/other.md","digest":"` + digest + `"}`}, } { t.Run(tc.name, func(t *testing.T) { - observed := map[string]any{"output": tc.output, "truncated": true, "state": "exited", "done": true} + observed := map[string]any{ + "output": tc.output, + "truncated": true, + "state": "exited", + "done": true, + } if tc.metadata != "" { observed["file_transfer"] = json.RawMessage(tc.metadata) } @@ -224,7 +134,9 @@ func TestFileTransferToolResultContentParts(t *testing.T) { process := ProcessRecord{ ExecutionSpec: processcmd.ForFileTransfer(processcmd.FileTransfer{ Direction: direction, - Target: processcmd.FileTarget{Memory: &processcmd.MemoryTarget{StoreID: uuid.New(), Path: "notes.md"}}, + Target: processcmd.FileTarget{ + Memory: &processcmd.MemoryTarget{StoreID: uuid.New(), Path: "notes.md"}, + }, }), State: ProcessStateExited, ExitCode: &exitCode, @@ -270,7 +182,9 @@ func TestFileTransferFailureResult(t *testing.T) { } { t.Run(tc.name, func(t *testing.T) { process := ProcessRecord{ - ExecutionSpec: processcmd.ForFileTransfer(processcmd.FileTransfer{Direction: processcmd.FileTransferUpload}), + ExecutionSpec: processcmd.ForFileTransfer( + processcmd.FileTransfer{Direction: processcmd.FileTransferUpload}, + ), State: tc.state, StateReasonCode: tc.reason, ExitCode: &tc.exitCode, diff --git a/internal/storage/executionstore/tool_overflow_integration_test.go b/internal/storage/executionstore/tool_overflow_integration_test.go index f5db1c48d..f213f99c6 100644 --- a/internal/storage/executionstore/tool_overflow_integration_test.go +++ b/internal/storage/executionstore/tool_overflow_integration_test.go @@ -18,7 +18,6 @@ import ( "github.com/omnara-ai/omnara/internal/publicid" "github.com/omnara-ai/omnara/internal/storage" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/toolcatalog" ) @@ -55,7 +54,10 @@ func (b *overflowBlobs) GetBlob(_ context.Context, key string) ([]byte, blobstor if !ok { return nil, blobstore.Metadata{}, blobstore.ErrNotFound } - return content, blobstore.Metadata{Digest: blobstore.ContentDigest(content), SizeBytes: int64(len(content))}, nil + return content, blobstore.Metadata{ + Digest: blobstore.ContentDigest(content), + SizeBytes: int64(len(content)), + }, nil } func (b *overflowBlobs) DeleteBlob(_ context.Context, key string) error { @@ -132,10 +134,11 @@ func TestToolOverflowConcurrentCompletions(t *testing.T) { workers.Add(1) go func() { defer workers.Done() - _, err := store.Execution().CompleteRuntimeToolCall(ctx, executionstore.CompleteRuntimeToolCallInput{ - ProjectID: testProjectID, AgentID: fixture.AgentID, ID: id, RuntimeLockID: fixture.Lock.ID, - Outcome: executionstore.ToolResultOutcomeSucceeded, ResultContentParts: parts, - }) + _, err := store.Execution(). + CompleteRuntimeToolCall(ctx, executionstore.CompleteRuntimeToolCallInput{ + ProjectID: testProjectID, AgentID: fixture.AgentID, ID: id, RuntimeLockID: fixture.Lock.ID, + Outcome: executionstore.ToolResultOutcomeSucceeded, ResultContentParts: parts, + }) results <- err }() } @@ -249,13 +252,18 @@ func TestToolOverflowCompletionAndReplay(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := store.Execution().CompleteToolCall(ctx, input); !errors.Is(err, storeerr.ErrIdempotencyConflict) { + if _, err := store.Execution().CompleteToolCall(ctx, input); !errors.Is( + err, + storeerr.ErrIdempotencyConflict, + ) { t.Fatalf("conflicting replay: %v", err) } if _, _, err := store.Artifacts().GetArtifactBlob(ctx, testProjectID, callID, id); err == nil { t.Fatal("cross-agent artifact access succeeded") } - for _, kind := range []string{toolcatalog.ToolTypeMCP, toolcatalog.ToolTypeCustom, toolcatalog.ToolTypeBuiltIn} { + for _, kind := range []string{toolcatalog.ToolTypeMCP, + toolcatalog.ToolTypeCustom, + toolcatalog.ToolTypeBuiltIn} { name := "read_process" if kind == toolcatalog.ToolTypeMCP { name = "mcp__docs__search" @@ -314,14 +322,17 @@ func TestToolOverflowExcludesProcessReadReplay(t *testing.T) { action, err := createProcessActionForTest(ctx, store, executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, AgentID: fixture.AgentID, ToolCallID: callID, RuntimeLockID: fixture.Lock.ID, }, executionstore.CreateProcessActionInput{ - ProcessID: process.ID, ActionKind: executionstore.ProcessActionKindRead, Payload: json.RawMessage(`{"max_bytes":65536}`), + ProcessID: process.ID, + ActionKind: executionstore.ProcessActionKindRead, + Payload: json.RawMessage(`{"max_bytes":65536}`), }) if err != nil { t.Fatal(err) } - _, found, err := store.Execution().AcceptDaemonProcessAction(ctx, executionstore.AcceptDaemonProcessActionInput{ - Authority: fixture.authority(), ProcessID: process.ID, ID: action.ID, - }) + _, found, err := store.Execution(). + AcceptDaemonProcessAction(ctx, executionstore.AcceptDaemonProcessActionInput{ + Authority: fixture.authority(), ProcessID: process.ID, ID: action.ID, + }) if err != nil || !found { t.Fatalf("accept read: found=%v err=%v", found, err) } @@ -385,7 +396,9 @@ func TestQuestionResponseSizeLimit(t *testing.T) { ctx := context.Background() pool := openIntegrationDB(t, ctx) seedMigratedDB(t, ctx, pool) - blobs := &overflowBlobs{beforeIO: func() { t.Fatal("question completion must not use blob storage") }} + blobs := &overflowBlobs{ + beforeIO: func() { t.Fatal("question completion must not use blob storage") }, + } store := newIntegrationStore(pool, storage.WithBlobStore(blobs)) user := mustCreateProjectOperatorUser(t, ctx, store, "bounded-question@example.com", "Question") fixture := newProcessDaemonFixtureInStore(t, ctx, store, user.ID, "bounded_question", time.Now()) @@ -398,7 +411,9 @@ func TestQuestionResponseSizeLimit(t *testing.T) { result, err := store.Execution().ExecuteToolCall(ctx, executionstore.ExecuteToolCallInput{ ProjectID: testProjectID, AgentID: fixture.AgentID, ToolCallID: callID, RuntimeLockID: fixture.Lock.ID, }, func(_ *executionstore.ToolCallReader) (executionstore.ToolCallCommand, error) { - return executionstore.CreateQuestionForToolCall(executionstore.CreateQuestionInteractionInput{Form: form}), nil + return executionstore.CreateQuestionForToolCall( + executionstore.CreateQuestionInteractionInput{Form: form}, + ), nil }) if err != nil { t.Fatal(err) @@ -408,12 +423,15 @@ func TestQuestionResponseSizeLimit(t *testing.T) { t.Fatalf("question command returned %T", result.CommandResult) } resultParts := func(answer string) json.RawMessage { - return mustTestRawJSON(t, []map[string]any{{"type": "structured_data", "value": map[string]any{ - "answers": []map[string]any{{ - "question_index": 0, "question": form.Questions[0].Prompt, "text": answer, - "selected_options": []map[string]any{{"option_index": 0, "label": "Yes"}}, - }}, - }}}) + return mustTestRawJSON( + t, + []map[string]any{{"type": "structured_data", "value": map[string]any{ + "answers": []map[string]any{{ + "question_index": 0, "question": form.Questions[0].Prompt, "text": answer, + "selected_options": []map[string]any{{"option_index": 0, "label": "Yes"}}, + }}, + }}}, + ) } answer := "answer" if test.size != 0 { @@ -430,23 +448,29 @@ func TestQuestionResponseSizeLimit(t *testing.T) { Actor: mustOmnaraActorParams(t, fixture.UserID), } if test.legacy { - q := dbsqlc.New(pool) actorID := fixture.omnaraActorID(t, ctx) - scope, key := "agent_interaction_response", interaction.ID.String() - response, err := q.InsertInteractionResponseAgentInput(ctx, dbsqlc.InsertInteractionResponseAgentInputParams{ - ProjectID: testProjectID, AgentID: fixture.AgentID, TargetInteractionID: interaction.ID, - ActorID: &actorID, - IdempotencyScope: &scope, InputIdempotencyKey: &key, - Metadata: json.RawMessage(`{}`), - }) - if err != nil { + responseID := uuid.New() + if _, + err := pool.Exec(ctx, + `INSERT INTO agent_inputs(id,agent_id,input_kind,state,target_interaction_id,actor_id,project_id,idempotency_scope,input_idempotency_key,metadata,queued_at,delivery_mode) +VALUES($1,$2,'interaction_response','received',$3,$4,$6,'agent_interaction_response',$5, +'{}',statement_timestamp(),'immediate')`, + responseID, + fixture.AgentID, + interaction.ID, + actorID, + interaction.ID.String(), + testProjectID); err != nil { t.Fatal(err) } - if _, err := q.ResolveAgentInteraction(ctx, dbsqlc.ResolveAgentInteractionParams{ - ProjectID: testProjectID, AgentID: fixture.AgentID, ID: interaction.ID, - Resolution: mustTestRawJSON(t, input.Resolution), - ResolvedByInputID: &response.ID, - }); err != nil { + if _, + err := pool.Exec(ctx, + `UPDATE agent_interactions SET state='resolved',resolution=$3,resolved_by_input_id=$4,resolved_at=statement_timestamp() WHERE agent_id=$1 AND id=$2`, + fixture.AgentID, + interaction.ID, + mustTestRawJSON(t, + input.Resolution), + responseID); err != nil { t.Fatal(err) } var blocks []struct{ Value json.RawMessage } @@ -465,15 +489,16 @@ func TestQuestionResponseSizeLimit(t *testing.T) { t.Fatalf("question resolution: state=%s, error=%v", resolved.State, err) } assertJSONRawEqual(t, resolved.Resolution, string(mustTestRawJSON(t, input.Resolution))) - if got := countAgentWakeups(t, ctx, store, fixture.AgentID); got != 1 { - t.Fatalf("wakeups=%d, want 1", got) + if got := countAgentWakeups(t, ctx, store, fixture.AgentID); got != 0 { + t.Fatalf("wakeups=%d, want 0 while owned", got) } call, err := store.Execution().GetToolCall(ctx, testProjectID, fixture.AgentID, callID) if err != nil || call.State != executionstore.ToolCallStateCompleted { t.Fatalf("question completion: state=%s, error=%v", call.State, err) } if test.wantFailed { - if call.Outcome != executionstore.ToolResultOutcomeFailed || len(call.ResultContentParts) > limit { + if call.Outcome != executionstore.ToolResultOutcomeFailed || + len(call.ResultContentParts) > limit { t.Fatalf("outcome=%s, bytes=%d", call.Outcome, len(call.ResultContentParts)) } if !strings.Contains(string(call.ResultContentParts), "request a shorter answer") { @@ -520,7 +545,9 @@ func TestToolOverflowRechecksArchiveDuringUpload(t *testing.T) { } _, err = store.Execution().CompleteToolCall(ctx, executionstore.CompleteToolCallInput{ ProjectID: testProjectID, AgentID: fixture.AgentID, ID: callID, - RuntimeLockID: fixture.Lock.ID, Outcome: executionstore.ToolResultOutcomeSucceeded, ResultContentParts: parts, + RuntimeLockID: fixture.Lock.ID, + Outcome: executionstore.ToolResultOutcomeSucceeded, + ResultContentParts: parts, }) if !errors.Is(err, storeerr.ErrStateTransitionConflict) { t.Fatalf("archived completion: %v", err) @@ -556,7 +583,10 @@ func TestToolOverflowStructuredReplayPreservesJSONEquality(t *testing.T) { "type": "structured_data", "value": map[string]any{"number": 1, "text": strings.Repeat("result", 10000)}, }} if name == "multiple large blocks" { - original = append(original, map[string]any{"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}) + original = append( + original, + map[string]any{"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}, + ) } if name == "text with metadata" || name == "hidden text block" { original = []map[string]any{ @@ -590,12 +620,15 @@ func TestToolOverflowStructuredReplayPreservesJSONEquality(t *testing.T) { Outcome: executionstore.ToolResultOutcomeSucceeded, ResultContentParts: parts, } for _, number := range numbers { - input.ResultContentParts = json.RawMessage(strings.Replace(string(parts), `"number":1`, `"number":`+number, 1)) + input.ResultContentParts = json.RawMessage( + strings.Replace(string(parts), `"number":1`, `"number":`+number, 1), + ) if runtime { - _, err = store.Execution().CompleteRuntimeToolCall(ctx, executionstore.CompleteRuntimeToolCallInput{ - ProjectID: input.ProjectID, AgentID: input.AgentID, ID: input.ID, RuntimeLockID: input.RuntimeLockID, - Outcome: input.Outcome, ResultContentParts: input.ResultContentParts, - }) + _, err = store.Execution(). + CompleteRuntimeToolCall(ctx, executionstore.CompleteRuntimeToolCallInput{ + ProjectID: input.ProjectID, AgentID: input.AgentID, ID: input.ID, RuntimeLockID: input.RuntimeLockID, + Outcome: input.Outcome, ResultContentParts: input.ResultContentParts, + }) } else { _, err = store.Execution().CompleteToolCall(ctx, input) } @@ -641,7 +674,13 @@ func TestToolOverflowExcludesBoundedAndControlTools(t *testing.T) { t.Fatal(err) } for _, name := range []string{ - "read_file", "search_files", "run_command", "read_process", "list_processes", "write_process", "stop_process", + "read_file", + "search_files", + "run_command", + "read_process", + "list_processes", + "write_process", + "stop_process", "upload_file", "download_file", "list_files", "write_file", "create_machine", "delete_machine", "inspect_machine", "set_interaction_handler", } { diff --git a/internal/storage/executionstore/tool_result_overflow.go b/internal/storage/executionstore/tool_result_overflow.go index 2cfd4cd35..1bc943ab9 100644 --- a/internal/storage/executionstore/tool_result_overflow.go +++ b/internal/storage/executionstore/tool_result_overflow.go @@ -12,7 +12,9 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/blobstore" "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/resourcemeta" "github.com/omnara-ai/omnara/internal/storage/artifactstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/textutil" @@ -141,20 +143,6 @@ func shouldOffloadToolResult(toolType, name string) bool { return false } -func (s *Store) loadToolResultForReplay( - ctx context.Context, - read *toolResultArtifactReadRequiredError, -) (context.Context, error) { - if s.artifacts == nil { - return ctx, errors.New("artifact storage is required for tool result replay") - } - content, _, err := s.artifacts.GetArtifactBlob(ctx, read.projectID, read.agentID, read.artifactID) - if err != nil { - return ctx, err - } - return context.WithValue(ctx, toolResultArtifactContentsKey{}, map[uuid.UUID][]byte{read.artifactID: content}), nil -} - func (s *Store) prepareToolResult( ctx context.Context, projectID, agentID, callID uuid.UUID, @@ -208,52 +196,17 @@ func (s *Store) prepareToolResult( if err != nil { return nil, err } - if artifact.Digest != blobstore.ContentDigest(overflow.content) || artifact.ContentType != overflow.contentType || - artifact.Filename != "tool-result" || artifact.SizeBytes == nil || + if artifact.Digest != blobstore.ContentDigest(overflow.content) || + artifact.ContentType != overflow.contentType || + artifact.Filename != "tool-result" || + artifact.SizeBytes == nil || *artifact.SizeBytes != int64(len(overflow.content)) { return nil, storeerr.ErrIdempotencyConflict } return overflow.contentParts(artifact.ID) } -func replayToolResultOverflow( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, callID uuid.UUID, - blocks []CreateContentBlockInput, - canonical json.RawMessage, -) (json.RawMessage, error) { - overflow, err := prepareToolResultOverflow(blocks, canonical) - if err != nil { - return nil, err - } - if overflow.content == nil { - return canonical, nil - } - record, err := qtx.GetArtifactByIdempotencyKey(ctx, dbsqlc.GetArtifactByIdempotencyKeyParams{ - ProjectID: projectID, - AgentID: agentID, - IdempotencyKey: overflowArtifactKey(callID, canonical), - }) - if err != nil { - return nil, err - } - return overflow.contentParts(record.ID) -} - -type toolResultArtifactContentsKey struct{} - -type toolResultArtifactReadRequiredError struct { - projectID uuid.UUID - agentID uuid.UUID - artifactID uuid.UUID -} - -func (e *toolResultArtifactReadRequiredError) Error() string { - return "tool result artifact must be read outside the transaction" -} - -func expandToolResultOverflow( +func (s *Store) expandToolResult( ctx context.Context, projectID, agentID uuid.UUID, parts json.RawMessage, @@ -263,7 +216,8 @@ func expandToolResultOverflow( return nil, err } if len(blocks) < 2 || blocks[0].BlockKind != ContentBlockKindStructuredData || - blocks[1].BlockKind != ContentBlockKindArtifact || !blocks[1].ExcludeFromModelContext { + blocks[1].BlockKind != ContentBlockKindArtifact || + !blocks[1].ExcludeFromModelContext { return parts, nil } var marker struct { @@ -271,7 +225,7 @@ func expandToolResultOverflow( Truncated bool `json:"truncated"` ContentType string `json:"content_type"` } - if err := json.Unmarshal(blocks[0].StructuredData, &marker); err != nil { + if err = json.Unmarshal(blocks[0].StructuredData, &marker); err != nil { return nil, err } id, err := publicid.Encode(publicid.KindArtifact, blocks[1].ArtifactID) @@ -281,41 +235,60 @@ func expandToolResultOverflow( if !marker.Truncated || marker.Path != toolcatalog.ArtifactVFSRoot+"/"+id { return parts, nil } - contents, _ := ctx.Value(toolResultArtifactContentsKey{}).(map[uuid.UUID][]byte) - if content, ok := contents[blocks[1].ArtifactID]; ok { - if marker.ContentType == "text/plain" { - return marshalToolResultContentBlocks(append([]CreateContentBlockInput{{ - BlockKind: ContentBlockKindText, TextContent: string(content), - }}, blocks[2:]...)) - } - return content, nil + if s.artifacts == nil { + return nil, errors.New("artifact storage is required for tool result replay") + } + content, _, err := s.artifacts.GetArtifactBlob(ctx, projectID, agentID, blocks[1].ArtifactID) + if err != nil { + return nil, err } - return nil, &toolResultArtifactReadRequiredError{ - projectID: projectID, agentID: agentID, artifactID: blocks[1].ArtifactID, + if marker.ContentType == "text/plain" { + return marshalToolResultContentBlocks( + append( + []CreateContentBlockInput{{BlockKind: ContentBlockKindText, TextContent: string(content)}}, + blocks[2:]...), + ) } + return content, nil } -func toolResultContentMatchesTx( +func (s *Store) replayToolContent( ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID, callID uuid.UUID, - stored json.RawMessage, - blocks []CreateContentBlockInput, - canonical json.RawMessage, -) (bool, error) { - if sameJSON(stored, canonical) { - return true, nil + projectID, agentID uuid.UUID, + original json.RawMessage, + err error, +) (json.RawMessage, error) { + var conflict *agentexecution.ContentConflictError + if !errors.As(err, &conflict) { + return nil, err + } + blocks := make([]CreateContentBlockInput, len(conflict.Stored)) + for i, part := range conflict.Stored { + metadata, decodeErr := resourcemeta.FromJSON(part.Metadata) + if decodeErr != nil { + return nil, decodeErr + } + blocks[i] = CreateContentBlockInput{ + Ordinal: part.Ordinal, + BlockKind: ContentBlockKind(part.Kind), + TextContent: part.Text, + StructuredData: part.Data, + ArtifactID: part.ArtifactID, + ToolCallID: part.ToolID, + ExcludeFromModelContext: part.Exclude, + Metadata: metadata, + } } - rewritten, err := replayToolResultOverflow(ctx, qtx, projectID, agentID, callID, blocks, canonical) - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return false, err + stored, encodeErr := marshalToolResultContentBlocks(blocks) + if encodeErr != nil { + return nil, encodeErr } - if err == nil && sameJSON(stored, rewritten) { - return true, nil + expanded, readErr := s.expandToolResult(ctx, projectID, agentID, stored) + if readErr != nil { + return nil, readErr } - expanded, err := expandToolResultOverflow(ctx, projectID, agentID, stored) - if err != nil { - return false, err + if !sameJSON(expanded, original) { + return nil, err } - return sameJSON(expanded, canonical), nil + return stored, nil } diff --git a/internal/storage/executionstore/tool_result_overflow_test.go b/internal/storage/executionstore/tool_result_overflow_test.go index 7ace32cf1..1bb262531 100644 --- a/internal/storage/executionstore/tool_result_overflow_test.go +++ b/internal/storage/executionstore/tool_result_overflow_test.go @@ -2,7 +2,6 @@ package executionstore import ( "bytes" - "context" "encoding/json" "errors" "regexp" @@ -53,18 +52,28 @@ func TestToolResultOverflowPreservesContentAndAttachments(t *testing.T) { })}, {name: "mixed visible and hidden bundle", hidden: true, parts: mustTestRawJSON(t, []map[string]any{ {"type": "text", "text": strings.Repeat("visible", 10000)}, - {"type": "text", "text": strings.Repeat("hidden", 10000), "metadata": map[string]string{"omnara_hidden": "true"}}, + {"type": "text", + "text": strings.Repeat("hidden", + 10000), + "metadata": map[string]string{"omnara_hidden": "true"}}, {"type": "media_ref", "artifact_id": imageID.String()}, })}, {name: "bundle with hidden media", hidden: true, parts: mustTestRawJSON(t, []map[string]any{ {"type": "structured_data", "value": strings.Repeat("x", 60000)}, - {"type": "media_ref", "artifact_id": imageID.String(), "metadata": map[string]string{"omnara_hidden": "true"}}, - })}, - {name: "plaintext with retained hidden blocks", contentType: "text/plain", parts: mustTestRawJSON(t, []map[string]any{ - {"type": "text", "text": strings.Repeat("visible", 10000)}, - {"type": "structured_data", "value": "hidden", "metadata": map[string]string{"omnara_hidden": "true"}}, - {"type": "media_ref", "artifact_id": imageID.String(), "metadata": map[string]string{"omnara_hidden": "true"}}, + {"type": "media_ref", + "artifact_id": imageID.String(), + "metadata": map[string]string{"omnara_hidden": "true"}}, })}, + {name: "plaintext with retained hidden blocks", + contentType: "text/plain", + parts: mustTestRawJSON(t, + []map[string]any{ + {"type": "text", "text": strings.Repeat("visible", 10000)}, + {"type": "structured_data", "value": "hidden", "metadata": map[string]string{"omnara_hidden": "true"}}, + {"type": "media_ref", + "artifact_id": imageID.String(), + "metadata": map[string]string{"omnara_hidden": "true"}}, + })}, {name: "text with metadata", contentType: "text/plain", parts: mustTestRawJSON(t, []map[string]any{ {"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}, {"type": "structured_data", "value": map[string]any{"status_code": 200}}, @@ -73,11 +82,14 @@ func TestToolResultOverflowPreservesContentAndAttachments(t *testing.T) { {"type": "structured_data", "value": map[string]any{"status_code": 200}}, {"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}, })}, - {name: "escaped metadata fits inline", contentType: "text/plain", parts: mustTestRawJSON(t, []map[string]any{ - {"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}, - {"type": "structured_data", "value": strings.Repeat("\t", 20900)}, - {"type": "media_ref", "artifact_id": imageID.String()}, - })}, + {name: "escaped metadata fits inline", + contentType: "text/plain", + parts: mustTestRawJSON(t, + []map[string]any{ + {"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}, + {"type": "structured_data", "value": strings.Repeat("\t", 20900)}, + {"type": "media_ref", "artifact_id": imageID.String()}, + })}, {name: "escaped metadata exceeds inline budget", parts: mustTestRawJSON(t, []map[string]any{ {"type": "text", "text": strings.Repeat("TARGET é\n", 8000)}, {"type": "structured_data", "value": strings.Repeat("\t", 21000)}, @@ -126,7 +138,8 @@ func TestToolResultOverflowPreservesContentAndAttachments(t *testing.T) { t.Fatalf("generated %s block hidden=%t, want %t", block.BlockKind, hidden, test.hidden) } } - if bytes.Contains(parts, []byte(imageID.String())) && !bytes.Contains(result, []byte(imageID.String())) { + if bytes.Contains(parts, []byte(imageID.String())) && + !bytes.Contains(result, []byte(imageID.String())) { t.Fatal("lost existing media") } repeated, err := prepareToolResultOverflow(inputBlocks, parts) @@ -137,12 +150,8 @@ func TestToolResultOverflowPreservesContentAndAttachments(t *testing.T) { if err != nil || !bytes.Equal(result, replay) { t.Fatal("rewrite is not deterministic") } - ctx := context.WithValue(t.Context(), toolResultArtifactContentsKey{}, map[uuid.UUID][]byte{artifactID: saved}) - expanded, err := expandToolResultOverflow(ctx, uuid.New(), uuid.New(), result) - if err != nil || !sameJSON(expanded, parts) { - t.Fatalf("round trip changed original parts: %v", err) - } - if test.name == "text with metadata" && len(regexp.MustCompile(`(?m)^TARGET`).FindAll(saved, -1)) != 8000 { + if test.name == "text with metadata" && + len(regexp.MustCompile(`(?m)^TARGET`).FindAll(saved, -1)) != 8000 { t.Fatal("artifact lost original text lines") } if !bytes.Contains(result, []byte("/artifacts/")) { @@ -189,13 +198,6 @@ func TestToolResultOverflowPreservesMediaAboveInlineLimit(t *testing.T) { t.Fatalf("media reference %d was changed or lost", index) } } - ctx := context.WithValue(t.Context(), toolResultArtifactContentsKey{}, map[uuid.UUID][]byte{ - artifactID: overflow.content, - }) - expanded, err := expandToolResultOverflow(ctx, uuid.New(), uuid.New(), result) - if err != nil || !sameJSON(expanded, parts) { - t.Fatalf("round trip changed original mixed-media content: %v", err) - } } func TestSmallAndMediaOnlyResultsStayUnchanged(t *testing.T) { diff --git a/internal/storage/executionstore/tool_rows.go b/internal/storage/executionstore/tool_rows.go index 92e35d072..ab914f516 100644 --- a/internal/storage/executionstore/tool_rows.go +++ b/internal/storage/executionstore/tool_rows.go @@ -9,68 +9,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" ) -func toolCallRecordFromInsertSQLC(row dbsqlc.InsertToolCallRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, row.CompletedAt, - ) -} - -func toolCallRecordFromCompleteSQLC(row dbsqlc.CompleteToolCallRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromRuntimeCompleteSQLC(row dbsqlc.CompleteRuntimeToolCallRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromWaitingCompleteSQLC( - row dbsqlc.CompleteWaitingBuiltInToolCallRow, -) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromFailRuntimeSQLC(row dbsqlc.FailRuntimeToolCallsRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromCustomCompleteSQLC(row dbsqlc.CompleteCustomToolCallRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - func toolCallRecordFromGetSQLC(row dbsqlc.GetToolCallRow) ToolCallRecord { return toolCallRecordFromSQLC( row.ID, row.ProjectID, row.AgentID, row.TurnID, @@ -91,16 +29,6 @@ func toolCallRecordFromProviderSQLC(row dbsqlc.GetToolCallByProviderCallRow) Too ) } -func toolCallRecordFromContextSQLC(row dbsqlc.ListToolCallsForModelContextRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, row.CompletedAt, - ) -} - func toolCallRecordFromRunnableSQLC( row dbsqlc.NextRunnableToolCallForModelOutputRow, ) ToolCallRecord { @@ -140,86 +68,6 @@ func toolCallRecordFromWatermarkSQLC( return record } -func toolCallRecordFromReadySQLC(row dbsqlc.MarkToolCallReadyRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, row.CompletedAt, - ) -} - -func toolCallRecordFromProcessCompleteSQLC( - row dbsqlc.CompleteToolCallFromProcessRow, -) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromQuestionInteractionCompleteSQLC( - row dbsqlc.CompleteToolCallFromQuestionInteractionRow, -) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromStartedProcessCompleteSQLC( - row dbsqlc.CompleteToolCallFromStartedProcessRow, -) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromProcessActionCompleteSQLC( - row dbsqlc.CompleteToolCallFromProcessActionRow, -) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromCancelSQLC(row dbsqlc.CancelNonTerminalToolCallsForAgentRow) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - -func toolCallRecordFromPermissionInteractionCompleteSQLC( - row dbsqlc.CompleteToolCallFromPermissionInteractionRow, -) ToolCallRecord { - return toolCallRecordFromSQLC( - row.ID, row.ProjectID, row.AgentID, row.TurnID, - row.SourceEventID, row.ModelCallContextID, row.ProviderCallID, - row.Name, row.Input, row.Type, - row.State, row.Outcome, row.RuntimeLockID, - row.ResultContentParts, row.CreatedAt, nil, - ) -} - func toolCallRecordFromSQLC( id uuid.UUID, projectID uuid.UUID, diff --git a/internal/storage/executionstore/tool_work_preparation.go b/internal/storage/executionstore/tool_work_preparation.go new file mode 100644 index 000000000..8fa7123e6 --- /dev/null +++ b/internal/storage/executionstore/tool_work_preparation.go @@ -0,0 +1,52 @@ +package executionstore + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type PreparedToolWork struct { + Context ModelCallContextRecord + Config AgentConfigRecord + FirstCall *ToolCallRecord +} + +func prepareToolWorkTx( + ctx context.Context, q *dbsqlc.Queries, projectID, agentID uuid.UUID, work ClaimedToolWork, +) (*PreparedToolWork, error) { + model, err := loadModelCallContextByID(ctx, q, projectID, agentID, work.ModelCallContextID) + if err != nil { + return nil, err + } + output, err := q.GetModelOutputByModelContext(ctx, dbsqlc.GetModelOutputByModelContextParams{ + ProjectID: projectID, AgentID: agentID, ModelCallContextID: work.ModelCallContextID, + }) + if err != nil { + return nil, fmt.Errorf("load tool work output: %w", err) + } + if model.State != ModelCallContextSucceeded || output.ID != work.ModelOutputID { + return nil, storeerr.ErrStateTransitionConflict + } + config, err := q.GetAgentConfig(ctx, dbsqlc.GetAgentConfigParams{ProjectID: projectID, ID: model.AgentConfigID}) + if err != nil { + return nil, fmt.Errorf("load tool work config: %w", err) + } + result := &PreparedToolWork{Context: model, Config: agentConfigRecordFromSQLC(config)} + first, err := q.NextRunnableToolCallForModelOutput(ctx, dbsqlc.NextRunnableToolCallForModelOutputParams{ + ProjectID: projectID, AgentID: agentID, ModelOutputID: work.ModelOutputID, + }) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return nil, fmt.Errorf("load first runnable tool: %w", err) + } + if err == nil { + record := toolCallRecordFromRunnableSQLC(first) + result.FirstCall = &record + } + return result, nil +} diff --git a/internal/storage/executionstore/turn_rows.go b/internal/storage/executionstore/turn_rows.go index e218a657c..8cca14b31 100644 --- a/internal/storage/executionstore/turn_rows.go +++ b/internal/storage/executionstore/turn_rows.go @@ -1,47 +1,9 @@ package executionstore import ( - "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) -func agentTurnRecordFromInsertSQLC(row dbsqlc.InsertAgentTurnRow) AgentTurnRecord { - return agentTurnRecordFromFields( - row.ID, - row.ProjectID, - row.AgentID, - row.TurnSequence, - row.LatestEventID, - row.LatestSemanticEventID, - ) -} - -func agentTurnRecordFromCurrentContinuableSQLC(row dbsqlc.CurrentContinuableAgentTurnRow) AgentTurnRecord { - return agentTurnRecordFromFields( - row.ID, - row.ProjectID, - row.AgentID, - row.TurnSequence, - row.LatestEventID, - row.LatestSemanticEventID, - ) -} - -func agentTurnRecordFromFields( - id, projectID, agentID uuid.UUID, - sequence int64, - latestEventID, latestSemanticEventID uuid.UUID, -) AgentTurnRecord { - return AgentTurnRecord{ - ID: id, - ProjectID: projectID, - AgentID: agentID, - TurnSequence: sequence, - LatestEventID: latestEventID, - LatestSemanticEventID: latestSemanticEventID, - } -} - func agentRuntimeLockRecordFromSQLC(row dbsqlc.AgentRuntimeLock) AgentRuntimeLockRecord { return AgentRuntimeLockRecord{ ID: row.ID, diff --git a/internal/storage/identitystore/principal_activity.go b/internal/storage/identitystore/principal_activity.go index 160e1ce25..81547440b 100644 --- a/internal/storage/identitystore/principal_activity.go +++ b/internal/storage/identitystore/principal_activity.go @@ -14,10 +14,10 @@ import ( func (s *Store) EnsureUserPrincipalStillActiveTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, principal PrincipalRecord, ) error { - return ensureUserPrincipalStillActive(ctx, s.q.WithTx(tx), principal) + return ensureUserPrincipalStillActive(ctx, dbsqlc.New(tx), principal) } func ensureUserPrincipalStillActive( diff --git a/internal/storage/integrationstore/inbox_integration_test.go b/internal/storage/integrationstore/inbox_integration_test.go index 46ad48d7e..4313aa93e 100644 --- a/internal/storage/integrationstore/inbox_integration_test.go +++ b/internal/storage/integrationstore/inbox_integration_test.go @@ -73,7 +73,7 @@ func newInboxFixture(t *testing.T) inboxFixture { }, }) require.NoError(t, err) - store := integrationstore.New(pool, executionstore.IntegrationAccess{}) + store := integrationstore.New(pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) integration, err := store.CreateIntegration(ctx, integrationstore.SaveIntegrationInput{ OrgID: ids.OrgID, ProjectID: ids.ProjectID, Name: "inbox-integration", IntegrationKind: integrationdefinition.SlackThread, @@ -631,7 +631,7 @@ func TestInboxJSONBNormalizedSizeBoundaryIsExplicitAndAtomic(t *testing.T) { func TestInboxSetupUpdateWaitsForAtomicAdmission(t *testing.T) { t.Parallel() f := newInboxFixture(t) - f.store = integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + f.store = integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) wrapper, err := secrets.NewLocalKeyWrapper("inbox-test", map[string][]byte{ "inbox-test": []byte("0123456789abcdef0123456789abcdef"), }) diff --git a/internal/storage/integrationstore/inbox_lease.go b/internal/storage/integrationstore/inbox_lease.go index 8a862c973..cc0bfad0c 100644 --- a/internal/storage/integrationstore/inbox_lease.go +++ b/internal/storage/integrationstore/inbox_lease.go @@ -22,14 +22,14 @@ import ( // caller-owned transaction. CheckLease must succeed after admission writes and // before commit. Roll back on any error. type IntegrationInboxLeaseTx struct { - tx pgx.Tx + tx dbsqlc.DBTX q *dbsqlc.Queries lease IntegrationInboxLease record IntegrationInboxRecord } func (s *Store) LockIntegrationInboxLeaseTx( - ctx context.Context, tx pgx.Tx, lease IntegrationInboxLease, additional ...uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, lease IntegrationInboxLease, additional ...uuid.UUID, ) (*IntegrationInboxLeaseTx, error) { // Include every integration from the frozen admission plan so no new integration gate is // acquired beneath the receipt lock. diff --git a/internal/storage/integrationstore/inbox_store.go b/internal/storage/integrationstore/inbox_store.go index eeef25144..e9957dacd 100644 --- a/internal/storage/integrationstore/inbox_store.go +++ b/internal/storage/integrationstore/inbox_store.go @@ -55,7 +55,7 @@ func (s *Store) AcceptIntegrationReceipt( } func (s *Store) enterInboxIntegration( - ctx context.Context, tx pgx.Tx, projectID, integrationID uuid.UUID, additional ...uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, projectID, integrationID uuid.UUID, additional ...uuid.UUID, ) error { integration, err := getIntegration(ctx, dbsqlc.New(tx), projectID, integrationID) if err != nil { diff --git a/internal/storage/integrationstore/inbox_terminal_integration_test.go b/internal/storage/integrationstore/inbox_terminal_integration_test.go index 6354dd4f9..804b7a878 100644 --- a/internal/storage/integrationstore/inbox_terminal_integration_test.go +++ b/internal/storage/integrationstore/inbox_terminal_integration_test.go @@ -28,7 +28,7 @@ func freezeInboxLaunchClaim( t, f.pool.QueryRow(f.ctx, `SELECT id FROM agent_profiles WHERE project_id=$1 LIMIT 1`, f.project).Scan(&profileID), ) - store := integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + store := integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) integration, err := store.UpdateIntegration( f.ctx, f.integrationID, diff --git a/internal/storage/integrationstore/integration_agent_conversation.go b/internal/storage/integrationstore/integration_agent_conversation.go index 3028d9da2..26d47d75e 100644 --- a/internal/storage/integrationstore/integration_agent_conversation.go +++ b/internal/storage/integrationstore/integration_agent_conversation.go @@ -19,7 +19,7 @@ const agentConversationStateKind = "agent_conversation" // AssignAgentIntegrationConversationTx requires the integration lifecycle gate and either the // agent lock or the agent's uncommitted insert in this transaction. func (s *Store) AssignAgentIntegrationConversationTx( - ctx context.Context, tx pgx.Tx, projectID, agentID, integrationID uuid.UUID, address ConversationAddress, + ctx context.Context, tx dbsqlc.DBTX, projectID, agentID, integrationID uuid.UUID, address ConversationAddress, ) error { if projectID == uuid.Nil || agentID == uuid.Nil || integrationID == uuid.Nil { return storeerr.InvalidRequest(errors.New("project, agent and integration are required")) @@ -89,7 +89,7 @@ type AgentIntegrationConversationTarget struct { func GetAgentIntegrationConversationTargetTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, agentID, integrationID uuid.UUID, ) (AgentIntegrationConversationTarget, bool, error) { q := dbsqlc.New(tx) diff --git a/internal/storage/integrationstore/integration_credential_authority.go b/internal/storage/integrationstore/integration_credential_authority.go index 9dfea2ca1..3bedee60b 100644 --- a/internal/storage/integrationstore/integration_credential_authority.go +++ b/internal/storage/integrationstore/integration_credential_authority.go @@ -16,7 +16,7 @@ import ( // gates. Call it before conversation or agent locks. Credential rotation preserves // callback authority; revoking project access does not. func LockIntegrationCredentialAccessTx( - ctx context.Context, tx pgx.Tx, projectID, integrationID uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, projectID, integrationID uuid.UUID, ) error { integration, err := getIntegration(ctx, dbsqlc.New(tx), projectID, integrationID) if errors.Is(err, storeerr.ErrNotFound) { @@ -36,7 +36,7 @@ func LockIntegrationCredentialAccessTx( } func lockAvailableIntegrationCredential( - ctx context.Context, tx pgx.Tx, orgID, projectID, secretID uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, orgID, projectID, secretID uuid.UUID, ) (dbsqlc.GetProjectAvailableSecretRow, error) { if _, err := secretops.LockReference(ctx, tx, orgID, secretID); err != nil { if errors.Is(err, pgx.ErrNoRows) { diff --git a/internal/storage/integrationstore/integration_launch_claims.go b/internal/storage/integrationstore/integration_launch_claims.go index f689f23ab..210f3f168 100644 --- a/internal/storage/integrationstore/integration_launch_claims.go +++ b/internal/storage/integrationstore/integration_launch_claims.go @@ -280,7 +280,7 @@ func inboxLaunchIdentities( return identities, nil } -func lockInboxLaunchConversations(ctx context.Context, tx pgx.Tx, projectID, integrationID uuid.UUID, +func lockInboxLaunchConversations(ctx context.Context, tx dbsqlc.DBTX, projectID, integrationID uuid.UUID, identities map[integrationLaunchIdentity]map[string]bool) error { addresses := map[ConversationAddress]bool{} for identity := range identities { diff --git a/internal/storage/integrationstore/integration_launch_claims_integration_test.go b/internal/storage/integrationstore/integration_launch_claims_integration_test.go index ab49fae2b..d7e520a7e 100644 --- a/internal/storage/integrationstore/integration_launch_claims_integration_test.go +++ b/internal/storage/integrationstore/integration_launch_claims_integration_test.go @@ -18,7 +18,7 @@ import ( func TestIntegrationLaunchClaimFreezesEntireRecipientSet(t *testing.T) { t.Parallel() f := newInboxFixture(t) - store := integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + store := integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) var profileID uuid.UUID require.NoError( t, diff --git a/internal/storage/integrationstore/integration_lifecycle.go b/internal/storage/integrationstore/integration_lifecycle.go index edc3f5157..4631a2b44 100644 --- a/internal/storage/integrationstore/integration_lifecycle.go +++ b/internal/storage/integrationstore/integration_lifecycle.go @@ -7,6 +7,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -73,11 +74,12 @@ func (s *Store) DeleteIntegration(ctx context.Context, orgID, projectID, id uuid } func (s *Store) deleteIntegrationOnce(ctx context.Context, orgID, projectID, id uuid.UUID) error { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return err } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() if err := lifecyclelock.EnterActiveProject(ctx, tx, orgID, projectID); err != nil { return err } @@ -102,7 +104,7 @@ func (s *Store) deleteIntegrationOnce(ctx context.Context, orgID, projectID, id for _, agentID := range ids { refs = append(refs, lifecyclelock.AgentRef{ProjectID: projectID, AgentID: agentID}) } - if err := lifecyclelock.Agents(ctx, tx, refs); err != nil { + if err := unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}); err != nil { return err } if err := q.DeleteIntegrationSubscriptions(ctx, dbsqlc.DeleteIntegrationSubscriptionsParams{ @@ -110,7 +112,7 @@ func (s *Store) deleteIntegrationOnce(ctx context.Context, orgID, projectID, id }); err != nil { return err } - if err := s.access.ClearIntegrationTargetsFromAgents(ctx, tx, projectID, id); err != nil { + if err := s.access.ClearIntegrationTargetsFromAgents(ctx, unit, projectID, id); err != nil { return err } if err := q.DeleteIntegrationTargets( @@ -131,5 +133,5 @@ func (s *Store) deleteIntegrationOnce(ctx context.Context, orgID, projectID, id if rows == 0 { return storeerr.ErrNotFound } - return tx.Commit(ctx) + return unit.Commit(ctx, "deleteIntegrationOnce") } diff --git a/internal/storage/integrationstore/integration_locks.go b/internal/storage/integrationstore/integration_locks.go index e4652f284..08d28054f 100644 --- a/internal/storage/integrationstore/integration_locks.go +++ b/internal/storage/integrationstore/integration_locks.go @@ -6,13 +6,12 @@ import ( "slices" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) func LockIntegrationsTx( - ctx context.Context, tx pgx.Tx, projectID uuid.UUID, referenced []uuid.UUID, additional ...uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, projectID uuid.UUID, referenced []uuid.UUID, additional ...uuid.UUID, ) error { // Compiled references remain readable after revocation; only additional integrations // supply live authority. @@ -20,7 +19,7 @@ func LockIntegrationsTx( } func lockIntegrationsTx( - ctx context.Context, tx pgx.Tx, projectID uuid.UUID, referenced, required []uuid.UUID, + ctx context.Context, tx dbsqlc.DBTX, projectID uuid.UUID, referenced, required []uuid.UUID, ) error { ids := make(map[uuid.UUID]bool) for _, id := range referenced { diff --git a/internal/storage/integrationstore/integration_profile_choices_integration_test.go b/internal/storage/integrationstore/integration_profile_choices_integration_test.go index 0f2fe069f..ab4381515 100644 --- a/internal/storage/integrationstore/integration_profile_choices_integration_test.go +++ b/internal/storage/integrationstore/integration_profile_choices_integration_test.go @@ -8,11 +8,12 @@ import ( "encoding/json" "errors" "fmt" - "github.com/omnara-ai/omnara/internal/testutil/integrationtest" "sync" "testing" "time" + "github.com/omnara-ai/omnara/internal/testutil/integrationtest" + "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/integrationdefinition" "github.com/omnara-ai/omnara/internal/storage/executionstore" @@ -34,7 +35,7 @@ type profileChoiceFixture struct { func newProfileChoiceFixture(t *testing.T) profileChoiceFixture { t.Helper() base := newInboxFixture(t) - base.store = integrationstore.New(base.pool, executionstore.IntegrationAccess{}) + base.store = integrationstore.New(base.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) var profileID, configID uuid.UUID require.NoError(t, base.pool.QueryRow(base.ctx, `SELECT id FROM agent_profiles WHERE project_id=$1`, base.project).Scan(&profileID)) @@ -682,7 +683,7 @@ func TestIntegrationProfileChoiceOwnerRecoveryAndReceiptRetention(t *testing.T) require.NotEqual(t, f.source.ClaimToken, recovered.ClaimToken) _, _, err = f.store.EnsureIntegrationProfileChoice(f.ctx, f.source.Lease(), f.input) require.ErrorIs(t, err, integrationstore.ErrIntegrationInboxLeaseLost) - restarted := integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + restarted := integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) resumed, created, err := restarted.EnsureIntegrationProfileChoice(f.ctx, recovered.Lease(), f.input) require.NoError(t, err) require.False(t, created) diff --git a/internal/storage/integrationstore/integration_routing.go b/internal/storage/integrationstore/integration_routing.go index 59d4e850f..b90bdcaf0 100644 --- a/internal/storage/integrationstore/integration_routing.go +++ b/internal/storage/integrationstore/integration_routing.go @@ -6,7 +6,6 @@ import ( "errors" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/integrationdefinition" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -54,7 +53,7 @@ func (s *Store) IntegrationRoutingCandidatesForInbox( func (s *Store) IntegrationRoutingCandidatesTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, integrationID uuid.UUID, conversation ConversationAddress, scopes []ConversationAddress, diff --git a/internal/storage/integrationstore/integration_runtime.go b/internal/storage/integrationstore/integration_runtime.go index f167fed70..e1c45ec24 100644 --- a/internal/storage/integrationstore/integration_runtime.go +++ b/internal/storage/integrationstore/integration_runtime.go @@ -162,7 +162,7 @@ func (s *Store) ClaimIntegrationRuntime( return claim, true, nil } -func (s *Store) lockRuntimeAuthority(ctx context.Context, tx pgx.Tx, revision IntegrationRuntimeRevision) error { +func (s *Store) lockRuntimeAuthority(ctx context.Context, tx dbsqlc.DBTX, revision IntegrationRuntimeRevision) error { if err := s.enterInboxIntegration(ctx, tx, revision.ProjectID, revision.IntegrationID); err != nil { return err } diff --git a/internal/storage/integrationstore/integration_setup.go b/internal/storage/integrationstore/integration_setup.go index cd3982045..4ef178f77 100644 --- a/internal/storage/integrationstore/integration_setup.go +++ b/internal/storage/integrationstore/integration_setup.go @@ -112,7 +112,7 @@ func IdempotencyScope(integration IntegrationRecord) string { func validateIntegrationCredential( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, input ConfigureIntegrationInput, ) error { credential, err := lockAvailableIntegrationCredential( diff --git a/internal/storage/integrationstore/integration_subscription_registration.go b/internal/storage/integrationstore/integration_subscription_registration.go index 8564c4a31..0b9230790 100644 --- a/internal/storage/integrationstore/integration_subscription_registration.go +++ b/internal/storage/integrationstore/integration_subscription_registration.go @@ -17,7 +17,7 @@ const maxIntegrationConversationAgents = 16 func PrepareIntegrationSubscriptionTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID uuid.UUID, attachment IntegrationSubscriptionAttachment, ) (RegisterIntegrationSubscriptionInput, error) { @@ -43,7 +43,7 @@ func PrepareIntegrationSubscriptionTx( func RegisterIntegrationSubscriptionTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, input RegisterIntegrationSubscriptionInput, ) (IntegrationSubscriptionRecord, error) { rows, err := RegisterIntegrationSubscriptionsTx(ctx, tx, []RegisterIntegrationSubscriptionInput{input}) @@ -57,7 +57,7 @@ func RegisterIntegrationSubscriptionTx( // conversation locks before locking or creating the agent, and hold them through registration. func RegisterIntegrationSubscriptionsTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, inputs []RegisterIntegrationSubscriptionInput, ) ([]IntegrationSubscriptionRecord, error) { if len(inputs) == 0 { diff --git a/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go b/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go index 5dd8162ee..9f1f2ee4f 100644 --- a/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go +++ b/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go @@ -7,6 +7,7 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/stretchr/testify/require" ) @@ -36,46 +37,28 @@ func TestClearDeletedIntegrationTargetsUsesAgentIdentity(t *testing.T) { f.exec(t, `UPDATE agents SET interaction_target_id=$2, interaction_handler_key='inbox-integration' WHERE id=$1`, current.ID, currentTarget) f.exec(t, `UPDATE agents SET interaction_target_id=$2, interaction_handler_key='inbox-integration', - interaction_auto_select=false, state='archived', archived_at=now() WHERE id=$1`, historical.ID, historicalTarget) + interaction_auto_select=false, state='archived', archived_at=now() WHERE id=$1`, + historical.ID, + historicalTarget) f.exec(t, `UPDATE agents SET interaction_target_id=$2, interaction_handler_key='other', interaction_auto_select=false WHERE id=$1`, moved.ID, retainedTarget) - f.exec(t, `INSERT INTO agents(org_id,project_id,state,name,current_config_id,created_at,updated_at) - SELECT org_id,project_id,'active','unrelated',current_config_id,now(),now() - FROM agents CROSS JOIN generate_series(1,10000) WHERE id=$1`, current.ID) - f.exec(t, "ANALYZE agents") - f.exec(t, "ANALYZE integration_targets") - parameters := map[string]any{"project_id": f.project, "integration_id": f.integrationID} - plan := explainInboxQueryFile(t, f, "integration_target_lifecycle.sql", - "ClearDeletedIntegrationTargetsFromAgents", parameters) - var inspected float64 - var walk func(inboxQueryPlan) - walk = func(node inboxQueryPlan) { - if node.Relation == "agents" && node.NodeType != "ModifyTable" && node.Loops > 0 { - inspected += (node.Rows + node.Filtered) * node.Loops - require.NotEqual(t, "Seq Scan", node.NodeType, "deleting an integration must not scan project agents") - t.Logf("agent lookup: %s using %s, rows=%g filtered=%g loops=%g", - node.NodeType, node.Index, node.Rows, node.Filtered, node.Loops) - } - for _, child := range node.Plans { - walk(child) - } - } - walk(plan) - require.Positive(t, inspected, "the plan must visit the selected agents") - require.LessOrEqual(t, inspected, float64(4), "only the integration's four target owners need inspection") - - query, args := bindInboxQueryFile(t, "integration_target_lifecycle.sql", - "ClearDeletedIntegrationTargetsFromAgents", parameters) - foreignQuery, foreignArgs := bindInboxQueryFile(t, "integration_target_lifecycle.sql", - "ClearDeletedIntegrationTargetsFromAgents", map[string]any{ - "project_id": uuid.New(), "integration_id": f.integrationID, - }) - result, err := f.pool.Exec(f.ctx, foreignQuery, foreignArgs...) - require.NoError(t, err) - require.Zero(t, result.RowsAffected()) - result, err = f.pool.Exec(f.ctx, query, args...) - require.NoError(t, err) - require.EqualValues(t, 2, result.RowsAffected(), "clear live and historical selections, including archived agents") + f.exec( + t, + `WITH identities AS MATERIALIZED (SELECT gen_random_uuid() AS id FROM generate_series(1,10000)), inserted AS ( +INSERT INTO agents(id,root_agent_id,org_id,project_id,state,name,current_config_id,created_at,updated_at) +SELECT identity.id,identity.id,agent.org_id,agent.project_id,'active','unrelated', +agent.current_config_id,statement_timestamp(),statement_timestamp() +FROM identities identity CROSS JOIN agents agent WHERE agent.id=$1 RETURNING id) +INSERT INTO agent_execution_state(agent_id,turn_continuable,incomplete_tools) SELECT id,false,false FROM inserted`, + current.ID, + ) + cell := agentexecution.NewCell("test", f.pool, nil, nil) + require.NoError(t, cell.Transact(f.ctx, func(u *agentexecution.Unit) error { + return u.ClearIntegrationTargets(f.ctx, uuid.New(), f.integrationID) + })) + require.NoError(t, cell.Transact(f.ctx, func(u *agentexecution.Unit) error { + return u.ClearIntegrationTargets(f.ctx, f.project, f.integrationID) + })) for _, expected := range []struct { agentID, targetID uuid.UUID handlerKey string @@ -102,7 +85,16 @@ func TestClearDeletedIntegrationTargetsUsesAgentIdentity(t *testing.T) { } require.Equal(t, expected.autoSelect, autoSelect, "deletion must preserve automatic selection mode") } - result, err = f.pool.Exec(f.ctx, query, args...) - require.NoError(t, err) - require.Zero(t, result.RowsAffected(), "clearing already removed selections is idempotent") + require.NoError(t, cell.Transact(f.ctx, func(u *agentexecution.Unit) error { + return u.ClearIntegrationTargets(f.ctx, f.project, f.integrationID) + })) + var selected int + require.NoError( + t, + f.pool.QueryRow(f.ctx, + `SELECT count(*) FROM agents a JOIN integration_targets t ON t.id=a.interaction_target_id WHERE t.integration_id=$1`, + f.integrationID). + Scan(&selected), + ) + require.Zero(t, selected) } diff --git a/internal/storage/integrationstore/integration_target_lookup.go b/internal/storage/integrationstore/integration_target_lookup.go index 1dcc16549..bdcb5f1bf 100644 --- a/internal/storage/integrationstore/integration_target_lookup.go +++ b/internal/storage/integrationstore/integration_target_lookup.go @@ -61,7 +61,7 @@ func (s *Store) GetAgentIntegrationLaunchOwner( func (s *Store) GetIntegrationTargetTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, id uuid.UUID, ) (IntegrationTargetRecord, error) { return getIntegrationTarget(ctx, dbsqlc.New(tx), projectID, id) diff --git a/internal/storage/integrationstore/integration_targets.go b/internal/storage/integrationstore/integration_targets.go index 1b5b502a3..89f2e3c68 100644 --- a/internal/storage/integrationstore/integration_targets.go +++ b/internal/storage/integrationstore/integration_targets.go @@ -43,7 +43,7 @@ type EnsureConversationTargetInput struct { func LockConversationTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, projectID, integrationID uuid.UUID, address ConversationAddress, ) error { @@ -60,7 +60,7 @@ func LockConversationTx( func (s *Store) EnsureConversationTargetTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, input EnsureConversationTargetInput, ) (IntegrationTargetRecord, error) { if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.IntegrationID == uuid.Nil { diff --git a/internal/storage/integrationstore/integration_targets_integration_test.go b/internal/storage/integrationstore/integration_targets_integration_test.go index db6a3ef48..618a12f9b 100644 --- a/internal/storage/integrationstore/integration_targets_integration_test.go +++ b/internal/storage/integrationstore/integration_targets_integration_test.go @@ -32,7 +32,7 @@ func TestIntegrationConversationLaunchOwnershipRemainsIndependentOfSubscriptions }, ) require.NoError(t, err) - store := integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + store := integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) createIntegration := func(name string) uuid.UUID { t.Helper() return f.addIntegration(t, name, integrationstore.IntegrationSettings(`{}`)).ID diff --git a/internal/storage/integrationstore/integrations.go b/internal/storage/integrationstore/integrations.go index 5f2cb1e5b..4b39050ff 100644 --- a/internal/storage/integrationstore/integrations.go +++ b/internal/storage/integrationstore/integrations.go @@ -177,7 +177,7 @@ func (s *Store) GetIntegrationByID(ctx context.Context, id uuid.UUID) (Integrati func (s *Store) GetIntegrationByIDTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, id uuid.UUID, ) (IntegrationRecord, error) { return getIntegrationByID(ctx, dbsqlc.New(tx), id) diff --git a/internal/storage/integrationstore/integrations_integration_test.go b/internal/storage/integrationstore/integrations_integration_test.go index 0fb4344d8..d1b9ec77c 100644 --- a/internal/storage/integrationstore/integrations_integration_test.go +++ b/internal/storage/integrationstore/integrations_integration_test.go @@ -3,9 +3,10 @@ package integrationstore_test import ( - "github.com/omnara-ai/omnara/internal/testutil/integrationtest" "testing" + "github.com/omnara-ai/omnara/internal/testutil/integrationtest" + "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/integrationdefinition" "github.com/omnara-ai/omnara/internal/storage/executionstore" @@ -17,7 +18,7 @@ import ( func TestIntegrationReferencesAndLifecycle(t *testing.T) { t.Parallel() f := newInboxFixture(t) - s := integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + s := integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) execution := executionstore.New(f.pool, executionstore.Config{}) var configID uuid.UUID require.NoError( @@ -69,7 +70,7 @@ func TestIntegrationReferencesAndLifecycle(t *testing.T) { func TestIntegrationIndependentCapabilitiesAndPagination(t *testing.T) { t.Parallel() f := newInboxFixture(t) - s := integrationstore.New(f.pool, executionstore.IntegrationAccess{}) + s := integrationstore.New(f.pool, integrationstore.Config{Access: executionstore.IntegrationAccess{}}) input := integrationstore.SaveIntegrationInput{ OrgID: f.org, ProjectID: f.project, diff --git a/internal/storage/integrationstore/scheduled_event.go b/internal/storage/integrationstore/scheduled_event.go index b271423fd..617368b18 100644 --- a/internal/storage/integrationstore/scheduled_event.go +++ b/internal/storage/integrationstore/scheduled_event.go @@ -31,7 +31,7 @@ type AcceptScheduledIntegrationEventInput struct { func (s *Store) AcceptScheduledIntegrationEventTx( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, input AcceptScheduledIntegrationEventInput, ) (IntegrationInboxRecord, bool, error) { payload, err := json.Marshal(input.Event) diff --git a/internal/storage/integrationstore/store.go b/internal/storage/integrationstore/store.go index 734f8e8d6..76031dd73 100644 --- a/internal/storage/integrationstore/store.go +++ b/internal/storage/integrationstore/store.go @@ -6,28 +6,36 @@ import ( "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" + "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" ) type Access interface { - ClearIntegrationTargetsFromAgents(context.Context, pgx.Tx, uuid.UUID, uuid.UUID) error + ClearIntegrationTargetsFromAgents(context.Context, *agentexecution.Unit, uuid.UUID, uuid.UUID) error +} + +type Config struct { + PostCommitPublisher notifications.PostCommitPublisher + Access Access } type Store struct { + cell *agentexecution.AgentCell pool *storeutil.Pool q *dbsqlc.Queries access Access } -func New(pool *pgxpool.Pool, access Access) *Store { +func New(pool *pgxpool.Pool, config Config) *Store { db := storeutil.WrapPool(pool) return &Store{ + cell: agentexecution.NewCell("primary", pool, config.PostCommitPublisher, nil), pool: db, q: dbsqlc.New(db), - access: access, + access: config.Access, } } diff --git a/internal/storage/internal/agentexecution/cell.go b/internal/storage/internal/agentexecution/cell.go new file mode 100644 index 000000000..c33dcbe27 --- /dev/null +++ b/internal/storage/internal/agentexecution/cell.go @@ -0,0 +1,63 @@ +package agentexecution + +import ( + "context" + "errors" + "time" + + "github.com/omnara-ai/omnara/internal/modelprotocol" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" +) + +type AgentCell struct { + id CellID + pool *storeutil.Pool + publisher notifications.PostCommitPublisher + retryBackoff func(int, string) time.Duration +} + +func NewCell( + id CellID, + pool *pgxpool.Pool, + publisher notifications.PostCommitPublisher, + retryBackoff func(int, string) time.Duration, +) *AgentCell { + if retryBackoff == nil { + retryBackoff = modelprotocol.RetryBackoff + } + return &AgentCell{ + id: id, + pool: storeutil.WrapPool(pool), + publisher: publisher, + retryBackoff: retryBackoff, + } +} + +func (c *AgentCell) ID() CellID { return c.id } + +func (c *AgentCell) Begin(ctx context.Context) (*Unit, error) { + if c.id == "" { + return nil, errors.New("agent cell identity is required") + } + tx, err := c.pool.BeginTx(ctx, pgx.TxOptions{}) + if err != nil { + return nil, err + } + return newUnit(c, tx), nil +} + +func (c *AgentCell) Transact(ctx context.Context, run func(*Unit) error) error { + u, err := c.Begin(ctx) + if err != nil { + return err + } + defer func() { _ = u.Rollback(context.WithoutCancel(ctx)) }() + if err := run(u); err != nil { + return err + } + return u.Commit(ctx, "agent execution") +} diff --git a/internal/storage/internal/agentexecution/commands_agents.go b/internal/storage/internal/agentexecution/commands_agents.go new file mode 100644 index 000000000..2b004da22 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_agents.go @@ -0,0 +1,254 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/agentmessage" + "github.com/omnara-ai/omnara/internal/interactionform" + "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type CreateAgentInput struct { + ID uuid.UUID + ProjectID uuid.UUID + ParentID uuid.UUID + ConfigID uuid.UUID + ProfileID uuid.UUID + Name string + SubagentKey string + IdempotencyKey string + ArchiveAfterIdleMinutes *int32 +} + +func (u *Unit) CreateAgent(ctx context.Context, input CreateAgentInput) (*Handle, error) { + if err := u.active(); err != nil { + return nil, err + } + if input.ID == uuid.Nil || input.ProjectID == uuid.Nil || input.ConfigID == uuid.Nil { + return nil, errors.New("agent, project and config are required") + } + root := input.ID + if input.ParentID != uuid.Nil { + parent := u.handles[input.ParentID] + if parent == nil || parent.route.ProjectID != input.ProjectID { + return nil, ErrLockPlan + } + current, err := executiondb.New(). + ReadExecutionAgentIdentity(ctx, + u.DB(), + executiondb.ReadExecutionAgentIdentityParams{ProjectID: input.ProjectID, + ID: input.ParentID}) + if err != nil { + return nil, err + } + if current.State != "active" { + return nil, storeerr.ErrStateTransitionConflict + } + root = parent.route.RootAgentID + } + _, err := executiondb.New().CreateExecutionAgent(ctx, u.DB(), executiondb.CreateExecutionAgentParams{ + ID: input.ID, + ProjectID: input.ProjectID, + RootID: root, + ParentID: nullableID(input.ParentID), + ConfigID: input.ConfigID, + ProfileID: nullableID( + input.ProfileID, + ), + Name: input.Name, + SubagentKey: input.SubagentKey, + IdempotencyKey: optionalText(input.IdempotencyKey), + ArchiveAfterIdleMinutes: input.ArchiveAfterIdleMinutes}) + if err != nil { + return nil, err + } + if err = u.CreatedAgent(input.ProjectID, input.ID, input.ParentID); err != nil { + u.aborted = true + return nil, err + } + h := u.handles[input.ID] + h.mutation = &executionMutation{head: ExecutionHead{AgentID: input.ID}, dirty: true, insertHead: true} + return h, nil +} + +func (h *Handle) SetInteractionTarget( + ctx context.Context, + id uuid.UUID, + handler string, + automatic bool, +) (bool, error) { + if _, err := h.Route(); err != nil { + return false, err + } + n, err := executiondb.New(). + WriteExecutionInteractionTarget(ctx, + h.unit.DB(), + executiondb.WriteExecutionInteractionTargetParams{ID: h.route.AgentID, + TargetID: nullableID(id), + HandlerKey: optionalText(handler), + AutoSelect: automatic}) + return n > 0, err +} + +func (h *Handle) Repair(ctx context.Context) (ExecutionSnapshot, error) { + snapshot, err := h.ReconstructExecution(ctx) + if err != nil { + h.unit.aborted = true + return ExecutionSnapshot{}, err + } + h.mutation = &executionMutation{head: snapshot.Head, loaded: &snapshot, dirty: true, insertHead: true} + return snapshot, nil +} + +func (h *Handle) parentEffect(ctx context.Context, kind, text, key string, interactionID uuid.UUID) error { + if h.route.RootAgentID == h.route.AgentID { + return nil + } + child, err := executiondb.New(). + ReadExecutionAgentIdentity(ctx, + h.unit.DB(), + executiondb.ReadExecutionAgentIdentityParams{ProjectID: h.route.ProjectID, + ID: h.route.AgentID}) + if err != nil { + return err + } + parent := h.unit.handles[valueOrZero(child.ParentAgentID)] + if parent == nil { + return ErrLockPlan + } + snapshot, err := parent.LoadExecution(ctx) + if err != nil { + return err + } + if snapshot.View.State == AgentArchived { + return nil + } + principal, err := publicid.Encode(publicid.KindAgent, h.route.AgentID) + if err != nil { + return err + } + tenant, err := publicid.Encode(publicid.KindOrganization, child.OrgID) + if err != nil { + return err + } + q := dbsqlc.New(h.unit.DB()) + actor, err := q.GetActorByIdentity( + ctx, + dbsqlc.GetActorByIdentityParams{ + ProjectID: h.route.ProjectID, + Provider: "omnara", + ProviderTenantID: &tenant, + ProviderUserID: principal, + }, + ) + var actorID uuid.UUID + if errors.Is(err, pgx.ErrNoRows) { + created, createErr := q.UpsertActorIdentity( + ctx, + dbsqlc.UpsertActorIdentityParams{ + ProjectID: h.route.ProjectID, + Provider: "omnara", + ProviderTenantID: &tenant, + ProviderUserID: principal, + Metadata: json.RawMessage(`{}`), + }, + ) + if createErr != nil { + return createErr + } + actorID = created.ID + } else if err != nil { + return err + } else { + actorID = actor.ID + } + metadata := map[string]any{ + "kind": kind, + "agent_id": principal, + "name": child.Name, + "key": child.SubagentKey, + } + if interactionID != uuid.Nil { + encoded, err := publicid.Encode(publicid.KindAgentInteraction, interactionID) + if err != nil { + return err + } + metadata["interaction_id"] = encoded + } + data, err := json.Marshal(map[string]any{"subagent_message": metadata}) + if err != nil { + return err + } + name := child.Name + if name == "" { + name = child.SubagentKey + } + if name == "" { + name = "agent" + } + body := agentmessage.Render(name, child.SubagentKey, principal, kind, text) + _, err = parent.ReceiveContent( + ctx, + ReceiveContentInput{ + ActorID: actorID, + DeliveryMode: "steering", + IdempotencyScope: "subagent_message", + IdempotencyKey: key, + Metadata: data, + Content: []Content{{Kind: "text", Text: body}}, + }, + ) + if err != nil { + return err + } + return err +} + +func (h *Handle) questionEffect(ctx context.Context, i Interaction, form interactionform.Form) error { + return h.parentEffect( + ctx, + "question", + interactionform.RenderQuestion(form), + "question:"+i.ID.String(), + i.ID, + ) +} + +func (u *Unit) ClearIntegrationTargets(ctx context.Context, projectID, integrationID uuid.UUID) error { + q := executiondb.New() + ids, err := q.ListExecutionIntegrationTargetAgents( + ctx, + u.DB(), + executiondb.ListExecutionIntegrationTargetAgentsParams{ + ProjectID: projectID, + IntegrationID: integrationID, + }, + ) + if err != nil { + return err + } + refs := make([]lifecyclelock.AgentRef, len(ids)) + for i, id := range ids { + refs[i] = lifecyclelock.AgentRef{ProjectID: projectID, AgentID: id} + } + if err = u.LockAgentRefs(ctx, refs, LifecycleAuthority{}); err != nil { + return err + } + return q.ClearExecutionIntegrationTargets( + ctx, + u.DB(), + executiondb.ClearExecutionIntegrationTargetsParams{ + ProjectID: projectID, + IntegrationID: integrationID, + AgentIds: ids, + }, + ) +} diff --git a/internal/storage/internal/agentexecution/commands_backlog_integration_test.go b/internal/storage/internal/agentexecution/commands_backlog_integration_test.go new file mode 100644 index 000000000..633a6d752 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_backlog_integration_test.go @@ -0,0 +1,145 @@ +//go:build integration + +package agentexecution_test + +import ( + "errors" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +type backlogRecord struct { + ID uuid.UUID + Mode string + Rank int64 + State string +} + +func readCommandBacklog(t *testing.T, u *agentexecution.Unit, a uuid.UUID) []backlogRecord { + t.Helper() + rows, err := u.DB().Query(t.Context(), + "SELECT id,delivery_mode,input_rank,state FROM agent_inputs WHERE agent_id=$1 ORDER BY input_rank,queued_at,id", + a) + require.NoError(t, err) + result, err := pgx.CollectRows(rows, pgx.RowToStructByPos[backlogRecord]) + require.NoError(t, err) + return result +} + +func TestExecutionBacklogOrderingAndReplay(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + rollback := errors.New("restore backlog") + for _, dense := range []bool{false, true} { + for _, position := range []string{"front", "back", "before", "after"} { + name := position + if dense { + name += "/rebalance" + } + t.Run(name, func(t *testing.T) { + err := u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + ids := make([]uuid.UUID, 3) + for i := range ids { + ids[i] = receiveCommand(t, h, "queued", "input").ID + if dense { + _, err := u.DB().Exec(t.Context(), + "UPDATE agent_inputs SET input_rank=$2 WHERE agent_id=$1 AND id=$3", a, i+1, ids[i]) + require.NoError(t, err) + } + } + change := agentexecution.BacklogChange{ + ID: ids[2], + DeliveryMode: "queued", + Position: position, + } + if position == "before" || position == "after" { + change.BeforeID = ids[0] + } + _, err := h.ChangeBacklog(t.Context(), change) + require.NoError(t, err) + expected := readCommandBacklog(t, u, a) + got := make([]uuid.UUID, len(expected)) + for i, row := range expected { + got[i] = row.ID + } + order := []uuid.UUID{ids[0], ids[1], ids[2]} + if position == "front" || position == "before" { + order = []uuid.UUID{ids[2], ids[0], ids[1]} + } + if position == "after" { + order = []uuid.UUID{ids[0], ids[2], ids[1]} + } + require.Equal(t, order, got) + assertCommandState(t, u, h) + changed, err := h.ChangeBacklog(t.Context(), change) + require.NoError(t, err) + require.False(t, changed) + require.Equal(t, expected, readCommandBacklog(t, u, a)) + return rollback + }) + require.ErrorIs(t, err, rollback) + }) + } + } +} + +func TestExecutionBacklogDeliveryReplay(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + first := receiveCommand(t, h, "queued", "first") + second := receiveCommand(t, h, "queued", "second") + third := receiveCommand(t, h, "steering", "third") + promote := agentexecution.BacklogChange{ID: second.ID, DeliveryMode: "steering"} + changed, err := h.ChangeBacklog(t.Context(), promote) + require.NoError(t, err) + require.True(t, changed) + rows := readCommandBacklog(t, u, a) + for _, row := range rows { + if row.ID == second.ID { + require.EqualValues(t, 2048, row.Rank) + } else { + require.EqualValues(t, 1024, row.Rank) + } + } + changed, err = h.ChangeBacklog(t.Context(), promote) + require.NoError(t, err) + require.False(t, changed) + require.Equal(t, rows, readCommandBacklog(t, u, a)) + assertCommandState(t, u, h) + changed, err = h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "queued"}, + ) + require.NoError(t, err) + require.True(t, changed) + assertCommandState(t, u, h) + err = u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "queued"}, + ) + return err + }) + require.ErrorIs(t, err, storeerr.ErrStateTransitionConflict) + _, err = h.AdmitInputs(t.Context()) + require.NoError(t, err) + changed, err = h.ChangeBacklog(t.Context(), promote) + require.NoError(t, err) + require.False(t, changed) + assertCommandState(t, u, h) + cancel := agentexecution.BacklogChange{ID: first.ID, DeliveryMode: "queued", Cancel: true} + changed, err = h.ChangeBacklog(t.Context(), cancel) + require.NoError(t, err) + require.True(t, changed) + err = u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.ChangeBacklog(t.Context(), cancel) + return err + }) + require.ErrorIs(t, err, storeerr.ErrStateTransitionConflict) + assertCommandState(t, u, h) +} diff --git a/internal/storage/internal/agentexecution/commands_cancel.go b/internal/storage/internal/agentexecution/commands_cancel.go new file mode 100644 index 000000000..91aae5f04 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_cancel.go @@ -0,0 +1,238 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/omnara-ai/omnara/internal/notifications" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type CancelInput struct { + ActorID uuid.UUID + Reason string + Message string + ForceRuntime bool +} + +type CanceledExecution struct { + Event ExecutionEvent + RuntimeID uuid.UUID + Changed bool +} + +func (h *Handle) Cancel(ctx context.Context, input CancelInput) (CanceledExecution, error) { + return executeCommand(ctx, h, func(m *executionMutation) (CanceledExecution, error) { + if input.Reason == "" || input.Message == "" { + return CanceledExecution{}, errors.New("cancellation reason and message are required") + } + q := executiondb.New() + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return CanceledExecution{}, err + } + live, err := q.ListExecutionLiveContexts( + ctx, + h.unit.DB(), + executiondb.ListExecutionLiveContextsParams{AgentID: h.route.AgentID}, + ) + if err != nil { + return CanceledExecution{}, err + } + if snapshot.View.State == AgentArchived && input.ForceRuntime { + view := snapshot.View + view.State = AgentActive + snapshot.Selection, err = SelectNext(view, time.Time{}) + if err != nil { + return CanceledExecution{}, err + } + } + updates := make(map[uuid.UUID]*notifications.AgentInteractionUpdate) + turn := uuid.Nil + if snapshot.Selection.TurnContinuable { + turn = m.head.CurrentTurnID + } + affected := turn != uuid.Nil || len(live) > 0 + result := CanceledExecution{} + runtime, err := q.ReadExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.ReadExecutionRuntimeParams{AgentID: h.route.AgentID}, + ) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return result, err + } + if err == nil { + if affected || input.ForceRuntime { + canceled, err := q.RequestExecutionCancel( + ctx, + h.unit.DB(), + executiondb.RequestExecutionCancelParams{AgentID: h.route.AgentID}, + ) + if err != nil { + return result, err + } + result.RuntimeID = canceled.ID + h.unit.Notifications(). + AddWorkerControlCancel(canceled.WorkerProcessID, h.route.AgentID, canceled.ID) + } else if runtime.CancelRequestedAt != nil { + result.RuntimeID = runtime.ID + h.unit.Notifications().AddWorkerControlCancel(runtime.WorkerProcessID, h.route.AgentID, runtime.ID) + } + } + for _, c := range live { + if runtime.ID == uuid.Nil || c.RuntimeLockID != runtime.ID { + return result, storeerr.ErrStateTransitionConflict + } + n, err := q.TeardownExecutionContext(ctx, h.unit.DB(), executiondb.TeardownExecutionContextParams{ + AgentID: h.route.AgentID, + ID: c.ID, + RuntimeID: runtime.ID, + State: "canceled", + ErrorKind: "canceled", + ErrorCode: input.Reason, + ErrorMessage: input.Message, ErrorDetails: json.RawMessage(`{}`)}) + if err != nil { + return result, err + } + if n != 1 { + return result, storeerr.ErrStateTransitionConflict + } + m.changed() + } + if !affected { + return result, nil + } + _, err = q.CancelExecutionInputs( + ctx, + h.unit.DB(), + executiondb.CancelExecutionInputsParams{AgentID: h.route.AgentID}, + ) + if err != nil { + return result, err + } + id, err := q.InsertExecutionStop( + ctx, + h.unit.DB(), + executiondb.InsertExecutionStopParams{ + AgentID: h.route.AgentID, + ActorID: nullableID(input.ActorID), + }, + ) + if err != nil { + return result, err + } + result.Event, err = h.appendEvent(ctx, turn, "agent_input", id, uuid.Nil, uuid.Nil, false) + if err != nil { + return result, err + } + n, err := q.ResolveExecutionInput( + ctx, + h.unit.DB(), + executiondb.ResolveExecutionInputParams{ + AgentID: h.route.AgentID, + ID: id, + EventID: &result.Event.ID, + }, + ) + if err != nil { + return result, err + } + if n != 1 { + return result, storeerr.ErrStateTransitionConflict + } + if turn != uuid.Nil { + if err = h.advanceTurn(ctx, result.Event, false); err != nil { + return result, err + } + ids, err := q.ListExecutionInteractions( + ctx, + h.unit.DB(), + executiondb.ListExecutionInteractionsParams{AgentID: h.route.AgentID, TurnID: turn}, + ) + if err != nil { + return result, err + } + resolution, err := json.Marshal(map[string]string{"reason": input.Reason}) + if err != nil { + return result, err + } + for _, interactionID := range ids { + row, err := q.ReadExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ReadExecutionInteractionParams{AgentID: h.route.AgentID, ID: interactionID}, + ) + if err != nil { + return result, err + } + n, err := q.ResolveExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ResolveExecutionInteractionParams{ + AgentID: h.route.AgentID, + ID: interactionID, + State: "canceled", + InputID: &id, + Resolution: resolution, + }, + ) + if err != nil { + return result, err + } + if n != 1 { + return result, storeerr.ErrStateTransitionConflict + } + i := interactionRecord(row) + i.State = "canceled" + updates[i.ToolID] = ¬ifications.AgentInteractionUpdate{ + ID: i.ID, + InteractionKind: i.Kind, + State: i.State, + } + } + if err = h.cancelTurnProcesses(ctx, turn); err != nil { + return result, err + } + calls, err := q.ListExecutionUnfinishedTools( + ctx, + h.unit.DB(), + executiondb.ListExecutionUnfinishedToolsParams{AgentID: h.route.AgentID, TurnID: &turn}, + ) + if err != nil { + return result, err + } + for _, callID := range calls { + call, err := h.tool(ctx, callID) + if err != nil { + return result, err + } + if _, + err = h.completeTool(ctx, + m, + call, + "canceled", + []Content{{Kind: "structured_data", + Data: json.RawMessage(`{"reason":"Agent canceled before this tool call completed."}`)}}, + uuid.Nil, updates[call.ID]); err != nil { + return result, err + } + } + } + m.head.StopSequence = result.Event.Sequence + m.head.AnsweredThroughSequence = max(m.head.AnsweredThroughSequence, result.Event.Sequence) + m.head.PendingToolOutputID = uuid.Nil + m.head.PendingOutputLimitID = uuid.Nil + m.head.PendingConfigInputID = uuid.Nil + m.head.PendingCheckpointID = uuid.Nil + m.changed() + result.Changed = true + return result, nil + }) +} diff --git a/internal/storage/internal/agentexecution/commands_children.go b/internal/storage/internal/agentexecution/commands_children.go new file mode 100644 index 000000000..0a6bbc7b2 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_children.go @@ -0,0 +1,127 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +func (h *Handle) child(ctx context.Context, route AgentRoute) (*Handle, error) { + child, err := h.unit.Agent(route) + if err != nil { + return nil, err + } + if route.ProjectID != h.route.ProjectID || route.RootAgentID != h.route.RootAgentID { + return nil, storeerr.ErrNotFound + } + identity, err := executiondb.New(). + ReadExecutionAgentIdentity(ctx, + h.unit.DB(), + executiondb.ReadExecutionAgentIdentityParams{ProjectID: route.ProjectID, + ID: route.AgentID}) + if err != nil { + return nil, err + } + if valueOrZero(identity.ParentAgentID) != h.route.AgentID { + return nil, storeerr.ErrNotFound + } + if identity.State != "active" { + return nil, storeerr.ErrStateTransitionConflict + } + return child, nil +} + +func (h *Handle) childCommand( + ctx context.Context, + route AgentRoute, + completion ToolCompletion, + apply func(*Handle) error, +) (CompletedTool, error) { + return executeCommand(ctx, h, func(_ *executionMutation) (CompletedTool, error) { + if err := h.FenceRuntime(ctx, completion.RuntimeLockID); err != nil { + return CompletedTool{}, err + } + call, err := h.tool(ctx, completion.ID) + if err != nil { + return CompletedTool{}, err + } + if call.State == "completed" { + return h.CompleteTool(ctx, completion) + } + if call.State != "ready" || call.Type != "built_in" { + return CompletedTool{}, storeerr.ErrInvalidToolCallDisposition + } + child, err := h.child(ctx, route) + if err != nil { + return CompletedTool{}, err + } + if err = apply(child); err != nil { + return CompletedTool{}, err + } + return h.CompleteTool(ctx, completion) + }) +} + +func (h *Handle) SendToChild( + ctx context.Context, + route AgentRoute, + actorID uuid.UUID, + message string, + completion ToolCompletion, +) (CompletedTool, error) { + return h.childCommand(ctx, route, completion, func(child *Handle) error { + if message == "" { + return errors.New("child message is required") + } + id, err := publicid.Encode(publicid.KindAgent, h.route.AgentID) + if err != nil { + return err + } + metadata, err := json.Marshal(map[string]any{"parent_message": map[string]string{"agent_id": id}}) + if err != nil { + return err + } + received, err := child.ReceiveContent( + ctx, + ReceiveContentInput{ + ActorID: actorID, + DeliveryMode: "steering", + IdempotencyScope: "subagent_message", + IdempotencyKey: "tool_call:" + completion.ID.String(), + Metadata: metadata, + Content: []Content{{Kind: "text", Text: message}}, + }, + ) + if err != nil { + return err + } + if received.Created { + _, err = child.SupersedeInteractions(ctx, received.ID) + } + return err + }) +} + +func (h *Handle) CancelChild( + ctx context.Context, + route AgentRoute, + actorID uuid.UUID, + completion ToolCompletion, +) (CompletedTool, error) { + return h.childCommand(ctx, route, completion, func(child *Handle) error { + _, err := child.Cancel( + ctx, + CancelInput{ + ActorID: actorID, + Reason: "agent_canceled", + Message: "The model call was canceled by the parent agent.", + }, + ) + return err + }) +} diff --git a/internal/storage/internal/agentexecution/commands_compaction.go b/internal/storage/internal/agentexecution/commands_compaction.go new file mode 100644 index 000000000..07a6929c0 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_compaction.go @@ -0,0 +1,347 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/modelprotocol" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type BeginCompactionInput struct { + Failure ModelFailure + SourceEnd int64 +} +type CompactionStart struct { + Model PreparedModel + Preempted bool + Admission InputAdmission +} + +func (h *Handle) BeginCompaction(ctx context.Context, input BeginCompactionInput) (CompactionStart, error) { + return executeCommand(ctx, h, func(m *executionMutation) (CompactionStart, error) { + if err := h.FenceRuntime(ctx, input.Failure.RuntimeLockID); err != nil { + return CompactionStart{}, err + } + q := executiondb.New() + source, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: input.Failure.ContextID}, + ) + if err != nil { + return CompactionStart{}, err + } + if input.Failure.ErrorKind == "" || input.Failure.ErrorMessage == "" { + return CompactionStart{}, errors.New("compaction requires failure evidence") + } + if input.SourceEnd <= 0 || input.SourceEnd > source.InputEventSequence || + source.OperationKind == "compaction" && + input.SourceEnd >= valueOrZero(source.SourceEventSequenceEnd) { + return CompactionStart{}, storeerr.ErrStateTransitionConflict + } + input.Failure.Recovery = RecoveryCompact + if source.OperationKind == "compaction" { + input.Failure.Recovery = RecoveryReduce + } + if source.State != "started" { + historical := input.Failure + historical.RuntimeLockID = source.RuntimeLockID + if _, err = h.finishAttempt( + ctx, m, source, ContextFailed, historical, input.Failure.Evidence, + ); err != nil { + return CompactionStart{}, err + } + id, err := q.LatestExecutionAttempt(ctx, h.unit.DB(), executiondb.LatestExecutionAttemptParams{ + AgentID: h.route.AgentID, Operation: "compaction", Watermark: source.InputEventSequence, + SourceEnd: &input.SourceEnd}) + if err == nil { + existing, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: id}, + ) + return CompactionStart{Model: attemptRecord(existing)}, err + } + if !errors.Is(err, pgx.ErrNoRows) { + return CompactionStart{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return CompactionStart{}, err + } + if snapshot.Head.CurrentTurnID != source.TurnID || + snapshot.Head.StopSequence > source.InputEventSequence || + snapshot.View.Turn.LatestSemantic.Sequence > source.InputEventSequence { + return CompactionStart{Preempted: true}, nil + } + } + covered, err := q.ExecutionCompactionBoundary( + ctx, + h.unit.DB(), + executiondb.ExecutionCompactionBoundaryParams{ + AgentID: h.route.AgentID, + Watermark: source.InputEventSequence, + }, + ) + if err != nil { + return CompactionStart{}, err + } + if input.SourceEnd <= covered { + return CompactionStart{}, storeerr.ErrAgentNotAdvanceable + } + if source.State == "started" { + if _, err = h.finishAttempt( + ctx, m, source, ContextFailed, input.Failure, input.Failure.Evidence, + ); err != nil { + return CompactionStart{}, err + } + } + + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return CompactionStart{}, err + } + if snapshot.View.Inputs.Steering { + admitted, err := h.admitInputs(ctx, m, AdmitAllSteering) + return CompactionStart{Preempted: true, Admission: admitted}, err + } + if snapshot.Head.CurrentTurnID != source.TurnID || + snapshot.Head.StopSequence > source.InputEventSequence || + snapshot.View.Turn.LatestSemantic.Sequence > source.InputEventSequence { + return CompactionStart{Preempted: true}, nil + } + id, err := q.LatestExecutionAttempt(ctx, h.unit.DB(), executiondb.LatestExecutionAttemptParams{ + AgentID: h.route.AgentID, + Operation: "compaction", + Watermark: source.InputEventSequence, + SourceEnd: &input.SourceEnd}) + if err == nil { + existing, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: id}, + ) + return CompactionStart{Model: attemptRecord(existing)}, err + } + if !errors.Is(err, pgx.ErrNoRows) { + return CompactionStart{}, err + } + governing := m.head.NormalContextID + if source.OperationKind == "compaction" { + governing = m.head.CompactionContextID + } + if governing != source.ID { + return CompactionStart{}, storeerr.ErrAgentNotAdvanceable + } + desired := attemptRecord(source) + desired.Context.ID = uuid.New() + desired.Context.Operation = OperationCompaction + desired.Context.Attempt = 1 + desired.Context.SourceEventSequenceEnd = input.SourceEnd + desired.Context.State = ContextStarted + desired.Context.Recovery = RecoveryNone + desired.Context.RetryAt = nil + desired.RuntimeLockID = input.Failure.RuntimeLockID + result, err := h.createAttempt(ctx, m, desired, snapshot) + return CompactionStart{Model: result}, err + }) +} + +type PublishCheckpointInput struct { + RuntimeLockID uuid.UUID + ContextID uuid.UUID + Summary string + Evidence ModelEvidence +} + +type PublishedCheckpoint struct { + ID uuid.UUID + Event ExecutionEvent + Created bool +} + +func (h *Handle) PublishCheckpoint( + ctx context.Context, + input PublishCheckpointInput, +) (PublishedCheckpoint, error) { + return executeCommand(ctx, h, func(m *executionMutation) (PublishedCheckpoint, error) { + if input.Summary == "" { + return PublishedCheckpoint{}, errors.New("checkpoint summary is required") + } + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return PublishedCheckpoint{}, err + } + q := executiondb.New() + source, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: input.ContextID}, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + if source.OperationKind != "compaction" || source.SourceEventSequenceEnd == nil { + return PublishedCheckpoint{}, storeerr.ErrStateTransitionConflict + } + failure := ModelFailure{ContextID: input.ContextID, RuntimeLockID: input.RuntimeLockID} + if source.State != "started" { + if _, err = h.finishAttempt(ctx, m, source, ContextSucceeded, failure, input.Evidence); err != nil { + return PublishedCheckpoint{}, err + } + row, err := q.FindExecutionCheckpoint( + ctx, + h.unit.DB(), + executiondb.FindExecutionCheckpointParams{AgentID: h.route.AgentID, ContextID: source.ID}, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + if row.Summary != input.Summary { + return PublishedCheckpoint{}, storeerr.ErrIdempotencyConflict + } + return PublishedCheckpoint{ID: row.ID, Event: ExecutionEvent{ID: row.EventID, TurnID: row.TurnID, + EventBoundary: EventBoundary{Sequence: row.Sequence, Time: row.CreatedAt}}}, nil + } + if source.RuntimeLockID != input.RuntimeLockID { + return PublishedCheckpoint{}, storeerr.ErrRuntimeLockInactive + } + covered, err := q.ExecutionCompactionBoundary( + ctx, + h.unit.DB(), + executiondb.ExecutionCompactionBoundaryParams{ + AgentID: h.route.AgentID, + Watermark: source.InputEventSequence, + }, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + end := *source.SourceEventSequenceEnd + if end <= covered { + return PublishedCheckpoint{}, storeerr.ErrStateTransitionConflict + } + valid, err := q.ValidateExecutionCheckpoint( + ctx, + h.unit.DB(), + executiondb.ValidateExecutionCheckpointParams{ + AgentID: h.route.AgentID, + StartSequence: covered + 1, + EndSequence: end, + }, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + if valid.Events != end-covered || valid.OpenAuthorities { + return PublishedCheckpoint{}, storeerr.ErrCheckpointBoundaryUnsafe + } + boundary, err := q.ExecutionPreparationBoundary( + ctx, + h.unit.DB(), + executiondb.ExecutionPreparationBoundaryParams{AgentID: h.route.AgentID}, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + opening, err := h.captureOpening(ctx, source.TurnID, boundary.Watermark, m.head.StopSequence) + if err != nil { + return PublishedCheckpoint{}, err + } + row, err := q.CreateExecutionCheckpoint( + ctx, + h.unit.DB(), + executiondb.CreateExecutionCheckpointParams{AgentID: h.route.AgentID, + SourceEnd: end, + ContextID: source.ID, + Summary: input.Summary, + OpeningIds: opening.InputIDs, + OpeningSequence: optionalSequence(opening.EventSequence)}, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + event, err := h.appendEvent( + ctx, + source.TurnID, + "context_checkpoint", + uuid.Nil, + uuid.Nil, + row.ID, + false, + ) + if err != nil { + return PublishedCheckpoint{}, err + } + if err = h.advanceTurn(ctx, event, false); err != nil { + return PublishedCheckpoint{}, err + } + if _, err = h.finishAttempt(ctx, m, source, ContextSucceeded, failure, input.Evidence); err != nil { + return PublishedCheckpoint{}, err + } + if source.TurnID == m.head.CurrentTurnID && event.Sequence > m.head.MaxNormalInputSequence && + event.Sequence >= m.head.StopSequence { + m.head.PendingCheckpointID = row.ID + } + m.changed() + return PublishedCheckpoint{ID: row.ID, Event: event, Created: true}, nil + }) +} + +type ResumeCompactionInput struct { + RuntimeLockID uuid.UUID + ParentID uuid.UUID + Watermark int64 + SourceEnd int64 +} + +func (h *Handle) ResumeCompaction(ctx context.Context, input ResumeCompactionInput) (CompactionStart, error) { + return executeCommand(ctx, h, func(_ *executionMutation) (CompactionStart, error) { + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return CompactionStart{}, err + } + parent, err := executiondb.New().ReadExecutionAttempt(ctx, h.unit.DB(), executiondb.ReadExecutionAttemptParams{ + AgentID: h.route.AgentID, ID: input.ParentID, + }) + if err != nil { + return CompactionStart{}, err + } + if parent.OperationKind != "normal" || parent.State != "failed" || + valueOrZero(parent.RecoveryKind) != "compact" || parent.InputEventSequence != input.Watermark { + return CompactionStart{}, storeerr.ErrAgentNotAdvanceable + } + var metadata modelenvelope.ProviderMetadata + if err := json.Unmarshal(parent.ProviderMetadata, &metadata); err != nil { + return CompactionStart{}, err + } + cost, _ := modelenvelope.ParseProviderReportedCostUSD(parent.ProviderReportedCostUsd) + return h.BeginCompaction(ctx, BeginCompactionInput{ + SourceEnd: input.SourceEnd, + Failure: ModelFailure{ + ContextID: parent.ID, RuntimeLockID: input.RuntimeLockID, + ErrorKind: modelprotocol.ErrorKind(parent.ErrorKind), ErrorCode: parent.ErrorCode, + ErrorMessage: parent.ErrorMessage, ErrorDetails: parent.ErrorDetails, + Evidence: ModelEvidence{ + APIFormat: modelprotocol.APIFormat(parent.ApiFormat), + APIVariant: modelprotocol.APIVariant(parent.ApiVariant), + RequestID: parent.ProviderRequestID, ResponseID: parent.ProviderResponseID, + Usage: modelenvelope.Usage{ + InputTokens: int(valueOrZero(parent.InputTokensTotal)), + UncachedInputTokens: int(valueOrZero(parent.UncachedInputTokens)), + CacheReadTokens: int(valueOrZero(parent.CacheReadInputTokens)), + CacheWriteTokens: int(valueOrZero(parent.CacheWriteInputTokens)), + OutputTokens: int(valueOrZero(parent.OutputTokensTotal)), + ReasoningTokens: int(valueOrZero(parent.ReasoningOutputTokens)), + }, + Cost: cost, Metadata: metadata, + }, + }, + }) + }) +} diff --git a/internal/storage/internal/agentexecution/commands_config.go b/internal/storage/internal/agentexecution/commands_config.go new file mode 100644 index 000000000..c268581fc --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_config.go @@ -0,0 +1,143 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ActivateConfigInput struct { + ID uuid.UUID + ConfigID uuid.UUID + ActorID uuid.UUID + IdempotencyKey string + Metadata json.RawMessage +} + +type ActivatedConfig struct { + InputID uuid.UUID + Event ExecutionEvent + Created bool +} + +func (h *Handle) captureOpening(ctx context.Context, turn uuid.UUID, watermark, stop int64) (Opening, error) { + rows, err := executiondb.New().CaptureExecutionOpening(ctx, h.unit.DB(), executiondb.CaptureExecutionOpeningParams{ + ProjectID: h.route.ProjectID, AgentID: h.route.AgentID, TurnID: turn, Watermark: watermark, StopSequence: stop}) + if err != nil { + return Opening{}, err + } + opening := Opening{InputIDs: make([]uuid.UUID, 0, len(rows))} + for i, row := range rows { + opening.InputIDs = append(opening.InputIDs, row.ID) + if i == 0 { + opening.EventSequence = row.Sequence + } + } + return opening, nil +} + +func (h *Handle) ActivateConfig(ctx context.Context, input ActivateConfigInput) (ActivatedConfig, error) { + return executeCommand(ctx, h, func(m *executionMutation) (ActivatedConfig, error) { + if input.ConfigID == uuid.Nil { + return ActivatedConfig{}, errors.New("config is required") + } + metadata, err := objectJSON(input.Metadata) + if err != nil { + return ActivatedConfig{}, err + } + q := executiondb.New() + existing, err := q.FindExecutionInput( + ctx, + h.unit.DB(), + executiondb.FindExecutionInputParams{AgentID: h.route.AgentID, + ID: input.ID, Scope: optionalText("agent_config_change"), Key: optionalText(input.IdempotencyKey)}, + ) + if err == nil { + if existing.InputKind != "config_change" || valueOrZero(existing.AgentConfigID) != input.ConfigID || + valueOrZero(existing.ActorID) != input.ActorID || !equalJSON(existing.Metadata, metadata) { + return ActivatedConfig{}, storeerr.ErrIdempotencyConflict + } + if existing.State != "resolved" { + return ActivatedConfig{}, storeerr.ErrStateTransitionConflict + } + event, err := q.FindExecutionInputEvent( + ctx, + h.unit.DB(), + executiondb.FindExecutionInputEventParams{AgentID: h.route.AgentID, InputID: &existing.ID}, + ) + return ActivatedConfig{InputID: existing.ID, Event: ExecutionEvent{ID: event.ID, TurnID: event.TurnID, + EventBoundary: EventBoundary{Sequence: event.Sequence, Time: event.CreatedAt}}}, err + } + if !errors.Is(err, pgx.ErrNoRows) { + return ActivatedConfig{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return ActivatedConfig{}, err + } + if input.ID == uuid.Nil { + input.ID = uuid.New() + } + row, err := q.ReceiveExecutionInput(ctx, h.unit.DB(), executiondb.ReceiveExecutionInputParams{ + ID: input.ID, AgentID: h.route.AgentID, ProjectID: h.route.ProjectID, Kind: "config_change", Mode: "immediate", + ConfigID: &input.ConfigID, ActorID: nullableID(input.ActorID), Scope: optionalText("agent_config_change"), + Key: optionalText(input.IdempotencyKey), Metadata: metadata}) + if err != nil { + return ActivatedConfig{}, err + } + turn := m.head.CurrentTurnID + newTurn := !snapshot.Selection.TurnContinuable || turn == uuid.Nil + if newTurn { + turn = uuid.New() + } + event, err := h.appendEvent(ctx, turn, "agent_input", row.ID, uuid.Nil, uuid.Nil, newTurn) + if err != nil { + return ActivatedConfig{}, err + } + opening, err := h.captureOpening(ctx, turn, event.Sequence, m.head.StopSequence) + if err != nil { + return ActivatedConfig{}, err + } + n, err := q.ResolveExecutionInput(ctx, h.unit.DB(), executiondb.ResolveExecutionInputParams{ + AgentID: h.route.AgentID, ID: row.ID, EventID: &event.ID, + OpeningIds: opening.InputIDs, OpeningSequence: optionalSequence(opening.EventSequence)}) + if err != nil { + return ActivatedConfig{}, err + } + if n != 1 { + return ActivatedConfig{}, storeerr.ErrStateTransitionConflict + } + n, err = q.ActivateExecutionConfig(ctx, h.unit.DB(), executiondb.ActivateExecutionConfigParams{ + AgentID: h.route.AgentID, ProjectID: h.route.ProjectID, ConfigID: input.ConfigID}) + if err != nil { + return ActivatedConfig{}, err + } + if n != 1 { + return ActivatedConfig{}, storeerr.ErrStateTransitionConflict + } + if newTurn { + if err = q.OpenExecutionTurn(ctx, + h.unit.DB(), + executiondb.OpenExecutionTurnParams{AgentID: h.route.AgentID, + ID: turn, + EventID: event.ID}); err != nil { + return ActivatedConfig{}, err + } + m.openTurn(turn) + } else { + if err = h.advanceTurn(ctx, event, true); err != nil { + return ActivatedConfig{}, err + } + if snapshot.View.Turn.FirstContentSequence > 0 && event.Sequence > m.head.MaxNormalInputSequence { + m.head.PendingConfigInputID = row.ID + } + } + m.changed() + return ActivatedConfig{InputID: row.ID, Event: event, Created: true}, nil + }) +} diff --git a/internal/storage/internal/agentexecution/commands_content.go b/internal/storage/internal/agentexecution/commands_content.go new file mode 100644 index 000000000..08d2ed362 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_content.go @@ -0,0 +1,224 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/omnara-ai/omnara/internal/jsoncanonical" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/dbsafe" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type Content struct { + Ordinal int32 + Kind string + Text string + Data json.RawMessage + ArtifactID uuid.UUID + ToolID uuid.UUID + Exclude bool + Metadata json.RawMessage +} + +type ContentConflictError struct{ Stored []Content } + +func (e *ContentConflictError) Error() string { return storeerr.ErrIdempotencyConflict.Error() } +func (e *ContentConflictError) Unwrap() error { return storeerr.ErrIdempotencyConflict } + +type ExecutionEvent struct { + ID uuid.UUID + TurnID uuid.UUID + EventBoundary +} + +func objectJSON(data json.RawMessage) (json.RawMessage, error) { + if len(data) == 0 { + return json.RawMessage(`{}`), nil + } + var object map[string]json.RawMessage + if err := json.Unmarshal(data, &object); err != nil { + return nil, err + } + if object == nil { + return nil, errors.New("JSON object is required") + } + return json.Marshal(object) +} + +func equalJSON(a, b json.RawMessage) bool { + return jsoncanonical.Equal(a, b) +} + +func optionalText(s string) *string { + if s == "" { + return nil + } + return &s +} +func optionalSequence(n int64) *int64 { + if n == 0 { + return nil + } + return &n +} + +func (h *Handle) appendEvent( + ctx context.Context, + turn uuid.UUID, + kind string, + input, output, checkpoint uuid.UUID, + opening bool, +) (ExecutionEvent, error) { + args := executiondb.AppendExecutionEventParams{ID: uuid.New(), + AgentID: h.route.AgentID, + ProjectID: h.route.ProjectID, + TurnID: nullableID(turn), + Kind: kind, + InputID: nullableID( + input, + ), OutputID: nullableID(output), CheckpointID: nullableID(checkpoint), Opening: opening} + if input != uuid.Nil { + args.IdempotencyKey = optionalText("agent_input:" + input.String()) + } + row, err := executiondb.New().AppendExecutionEvent(ctx, h.unit.DB(), args) + if err != nil { + return ExecutionEvent{}, err + } + h.unit.Notifications().AddAgentEvent(h.route.AgentID, row.Sequence, kind) + return ExecutionEvent{ + ID: row.ID, + TurnID: turn, + EventBoundary: EventBoundary{Sequence: row.Sequence, Time: row.CreatedAt}, + }, nil +} + +func (h *Handle) writeContent(ctx context.Context, kind string, id uuid.UUID, parts []Content) error { + for n, part := range parts { + ordinal := int32(n) + if kind == "model_output" { + ordinal = part.Ordinal + } + if err := dbsafe.Text(part.Text); err != nil { + return err + } + metadata, err := objectJSON(part.Metadata) + if err != nil { + return err + } + if err := dbsafe.JSONStrings(metadata); err != nil { + return err + } + args := executiondb.InsertExecutionContentParams{AgentID: h.route.AgentID, + ProjectID: h.route.ProjectID, + OwnerKind: kind, + Ordinal: ordinal, + Kind: part.Kind, + ArtifactID: nullableID(part.ArtifactID), + ToolID: nullableID(part.ToolID), + Exclude: part.Exclude, + Metadata: metadata} + switch kind { + case "agent_input": + args.InputID = &id + case "model_output": + args.OutputID = &id + case "tool_call_result": + args.ResultID = &id + } + if part.Kind == "text" || part.Kind == "reasoning" || part.Kind == "error" { + args.Text = &part.Text + } + if len(part.Data) > 0 { + if !json.Valid(part.Data) { + return errors.New("invalid structured content") + } + if err := dbsafe.JSONStrings(part.Data); err != nil { + return err + } + args.Data = &part.Data + } + if err := executiondb.New().InsertExecutionContent(ctx, h.unit.DB(), args); err != nil { + return err + } + } + return nil +} + +func (h *Handle) contentMatches(ctx context.Context, input, output uuid.UUID, parts []Content) error { + return h.matchContent(ctx, executiondb.ReadExecutionContentParams{ + AgentID: h.route.AgentID, InputID: nullableID(input), OutputID: nullableID(output)}, parts) +} + +func (h *Handle) matchContent( + ctx context.Context, + args executiondb.ReadExecutionContentParams, + parts []Content, +) error { + rows, err := executiondb.New().ReadExecutionContent(ctx, h.unit.DB(), args) + if err != nil { + return err + } + conflict := &ContentConflictError{Stored: make([]Content, len(rows))} + for i, row := range rows { + conflict.Stored[i] = Content{ + Ordinal: row.Ordinal, + Kind: row.BlockKind, + Text: valueOrZero(row.TextContent), + Data: valueOrZero(row.StructuredData), + ArtifactID: valueOrZero(row.ArtifactID), + ToolID: valueOrZero(row.ToolCallID), + Exclude: row.ExcludeFromModelContext, + Metadata: row.Metadata, + } + } + if len(rows) != len(parts) { + return conflict + } + for i, row := range rows { + part := parts[i] + metadata, err := objectJSON(part.Metadata) + if err != nil { + return err + } + ordinal := int32(i) + if args.OutputID != nil { + ordinal = part.Ordinal + } + if row.Ordinal != ordinal || row.BlockKind != part.Kind || valueOrZero(row.TextContent) != part.Text || + valueOrZero(row.ArtifactID) != part.ArtifactID || + valueOrZero(row.ToolCallID) != part.ToolID || + row.ExcludeFromModelContext != part.Exclude || + !equalJSON(row.Metadata, metadata) || + ((len(part.Data) > 0 || + row.StructuredData != nil) && + !equalJSON(valueOrZero(row.StructuredData), + part.Data)) { + return conflict + } + } + return nil +} + +func (h *Handle) advanceTurn(ctx context.Context, event ExecutionEvent, semantic bool) error { + n, err := executiondb.New().AdvanceExecutionTurn(ctx, h.unit.DB(), executiondb.AdvanceExecutionTurnParams{ + AgentID: h.route.AgentID, TurnID: event.TurnID, EventID: event.ID, Semantic: semantic}) + if err == nil && n != 1 { + return ErrInvalidState + } + return err +} + +func (m *executionMutation) openTurn(id uuid.UUID) { + m.head.CurrentTurnID = id + m.head.NormalContextID = uuid.Nil + m.head.CompactionContextID = uuid.Nil + m.head.PendingToolOutputID = uuid.Nil + m.head.PendingOutputLimitID = uuid.Nil + m.head.PendingConfigInputID = uuid.Nil + m.head.PendingCheckpointID = uuid.Nil + m.changed() +} diff --git a/internal/storage/internal/agentexecution/commands_custody_integration_test.go b/internal/storage/internal/agentexecution/commands_custody_integration_test.go new file mode 100644 index 000000000..a529b5e34 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_custody_integration_test.go @@ -0,0 +1,217 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/interactionform" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +func toolCommand( + t *testing.T, + h *agentexecution.Handle, + lease uuid.UUID, + kind string, +) agentexecution.ToolRef { + t.Helper() + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = []modelenvelope.ResponsePart{ + {Type: "tool_call", ProviderCallID: "call", ToolName: "tool", ToolInput: json.RawMessage(`{}`)}, + } + output, err := h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ + RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: response, + Tools: []agentexecution.ToolProposal{{ProviderCallID: "call", Type: kind}}, + }, + ) + require.NoError(t, err) + return agentexecution.ToolRef{ID: output.Tools[0].ID, RuntimeLockID: lease} +} + +func TestExecutionToolCustody(t *testing.T) { + for _, kind := range []string{"built_in", "mcp", "custom"} { + t.Run(kind, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + ref := toolCommand(t, h, lease, kind) + assertCommandState(t, u, h) + changed, err := h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + require.True(t, changed) + assertCommandState(t, u, h) + changed, err = h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + require.False(t, changed) + if kind != "custom" { + changed, err = h.RunTool(t.Context(), ref) + require.NoError(t, err) + require.True(t, changed) + assertCommandState(t, u, h) + require.ErrorIs( + t, + u.Savepoint( + t.Context(), + func(_ *agentexecution.Unit) error { _, err := h.RunTool(t.Context(), ref); return err }, + ), + storeerr.ErrToolCallInProgress, + ) + changed, err = h.RequeueTool(t.Context(), ref) + require.NoError(t, err) + require.True(t, changed) + _, err = h.RunTool(t.Context(), ref) + require.NoError(t, err) + } else { + work, err := h.Advance(t.Context(), lease, true, true) + require.NoError(t, err) + require.True(t, work.Released) + } + content := []agentexecution.Content{{Kind: "text", Text: "result"}} + complete := func() (agentexecution.CompletedTool, error) { + if kind == "custom" { + return h.CompleteCustomTool(t.Context(), ref.ID, "succeeded", content) + } + return h.CompleteTool( + t.Context(), + agentexecution.ToolCompletion{ToolRef: ref, Outcome: "succeeded", Content: content}, + ) + } + result, err := complete() + require.NoError(t, err) + require.True(t, result.Changed) + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.WorkModel, snapshot.Selection.Work) + replay, err := complete() + require.NoError(t, err) + require.False(t, replay.Changed) + require.Equal(t, result.ResultID, replay.ResultID) + require.ErrorIs( + t, + u.Savepoint( + t.Context(), + func(_ *agentexecution.Unit) error { content[0].Text = "other"; _, err := complete(); return err }, + ), + storeerr.ErrIdempotencyConflict, + ) + require.NoError(t, u.Commit(t.Context(), "tool custody")) + }) + } +} + +func TestExecutionQuestionCustody(t *testing.T) { + for _, cancel := range []bool{false, true} { + t.Run(map[bool]string{false: "resolve", true: "supersede"}[cancel], func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + ref := toolCommand(t, h, lease, "built_in") + _, err := h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + form := interactionform.Form{ + Title: "Question", + Questions: []interactionform.Question{ + {Prompt: "Continue?", Options: []interactionform.Option{{Label: "yes"}}}, + }, + } + opened, err := h.OpenInteraction( + t.Context(), + agentexecution.OpenInteractionInput{ToolRef: ref, Question: &form}, + ) + require.NoError(t, err) + require.True(t, opened.Created) + assertCommandState(t, u, h) + replay, err := h.OpenInteraction( + t.Context(), + agentexecution.OpenInteractionInput{ToolRef: ref, Question: &form}, + ) + require.NoError(t, err) + require.False(t, replay.Created) + work, err := h.Advance(t.Context(), lease, true, true) + require.NoError(t, err) + require.True(t, work.Released) + assertCommandState(t, u, h) + if cancel { + input := receiveCommand(t, h, "steering", "cancel question") + ids, err := h.SupersedeInteractions(t.Context(), input.ID) + require.NoError(t, err) + require.Equal(t, []uuid.UUID{opened.ID}, ids) + } else { + input := agentexecution.ResolveInteractionInput{ID: opened.ID, + Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{OptionIndices: []int{0}}}}} + resolved, err := h.ResolveInteraction(t.Context(), input) + require.NoError(t, err) + require.Equal(t, "resolved", resolved.State) + replay, err := h.ResolveInteraction(t.Context(), input) + require.NoError(t, err) + require.Equal(t, resolved.InputID, replay.InputID) + } + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "question custody")) + }) + } +} + +func TestExecutionRuntimeRecoveryCommands(t *testing.T) { + for _, reap := range []bool{false, true} { + t.Run(map[bool]string{false: "release", true: "reap"}[reap], func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + if reap { + _, err = u.DB(). + Exec(t.Context(), + `UPDATE agent_runtime_locks SET started_at=statement_timestamp()-interval '3 minutes',renewed_at=statement_timestamp()-interval '2 minutes',lease_expires_at=statement_timestamp()-interval '1 minute' WHERE id=$1`, + lease) + require.NoError(t, err) + changed, err := h.Reap(t.Context(), lease) + require.NoError(t, err) + require.True(t, changed) + } else { + require.NoError(t, h.Release(t.Context(), lease)) + } + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.RecoveryRetry, snapshot.View.NormalContext.Recovery) + require.Equal(t, prepared.Context.ID, snapshot.View.NormalContext.ID) + require.Equal(t, agentexecution.WaitModelDeadline, snapshot.Selection.Wait) + require.NoError(t, u.Commit(t.Context(), "runtime recovery")) + }) + } +} + +func TestExecutionCancelCommand(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + toolCommand(t, h, lease, "built_in") + receiveCommand(t, h, "queued", "preserved") + receiveCommand(t, h, "steering", "canceled") + input := agentexecution.CancelInput{ + Reason: "agent_canceled", + Message: "The model call was canceled by an explicit agent cancellation.", + } + result, err := h.Cancel(t.Context(), input) + require.NoError(t, err) + require.True(t, result.Changed) + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.AdmitOneQueued, snapshot.Selection.Admission) + replay, err := h.Cancel(t.Context(), input) + require.NoError(t, err) + require.False(t, replay.Changed) + require.NoError(t, h.Release(t.Context(), lease)) + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "cancel")) +} diff --git a/internal/storage/internal/agentexecution/commands_failure.go b/internal/storage/internal/agentexecution/commands_failure.go new file mode 100644 index 000000000..ff450e34f --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_failure.go @@ -0,0 +1,265 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "math" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/modelprotocol" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ModelEvidence struct { + APIFormat modelprotocol.APIFormat + APIVariant modelprotocol.APIVariant + RequestID string + ResponseID string + Usage modelenvelope.Usage + Cost modelenvelope.ProviderReportedCostUSD + Metadata modelenvelope.ProviderMetadata +} + +type ModelFailure struct { + ContextID uuid.UUID + RuntimeLockID uuid.UUID + Recovery RecoveryKind + RetryDelay time.Duration + ErrorKind modelprotocol.ErrorKind + ErrorCode string + ErrorMessage string + ErrorDetails json.RawMessage + ServedModel string + Evidence ModelEvidence +} + +func normalizedUsage(usage modelenvelope.Usage) modelenvelope.Usage { + usage = modelenvelope.NormalizeUsage(usage) + for _, n := range []int{usage.InputTokens, + usage.UncachedInputTokens, + usage.CacheReadTokens, + usage.CacheWriteTokens, + usage.OutputTokens, + usage.ReasoningTokens} { + if n > math.MaxInt32 { + return modelenvelope.Usage{} + } + } + return usage +} + +func tokenCount(n int) *int32 { + if n <= 0 { + return nil + } + value := int32(n) + return &value +} + +func evidenceParams(e ModelEvidence) (executiondb.FinishExecutionAttemptParams, error) { + if err := modelenvelope.ValidateProviderReportedCostUSD(e.Cost); err != nil { + return executiondb.FinishExecutionAttemptParams{}, err + } + if (e.APIFormat == "") != (e.APIVariant == "") || + e.APIFormat == "" && + (e.RequestID != "" || e.ResponseID != "" || e.Cost != "" || e.Usage != (modelenvelope.Usage{})) { + return executiondb.FinishExecutionAttemptParams{}, errors.New( + "provider evidence requires API identity", + ) + } + metadata, err := json.Marshal(e.Metadata) + if err != nil { + return executiondb.FinishExecutionAttemptParams{}, err + } + usage := normalizedUsage(e.Usage) + return executiondb.FinishExecutionAttemptParams{ + ApiFormat: string(e.APIFormat), + ApiVariant: string(e.APIVariant), + RequestID: e.RequestID, + ResponseID: e.ResponseID, + InputTokens: tokenCount(usage.InputTokens), + UncachedTokens: tokenCount(usage.UncachedInputTokens), + CacheReadTokens: tokenCount( + usage.CacheReadTokens, + ), + CacheWriteTokens: tokenCount(usage.CacheWriteTokens), + OutputTokens: tokenCount(usage.OutputTokens), + ReasoningTokens: tokenCount(usage.ReasoningTokens), + Cost: optionalText(string(e.Cost)), + ProviderMetadata: metadata, + ErrorDetails: json.RawMessage(`{}`), + }, nil +} + +func (h *Handle) finishAttempt( + ctx context.Context, + m *executionMutation, + row executiondb.ReadExecutionAttemptRow, + state ContextState, + failure ModelFailure, + evidence ModelEvidence, +) (bool, error) { + args, err := evidenceParams(evidence) + if err != nil { + return false, err + } + details, err := objectJSON(failure.ErrorDetails) + if err != nil { + return false, err + } + if row.RuntimeLockID != failure.RuntimeLockID { + return false, storeerr.ErrRuntimeLockInactive + } + storedCost, _ := modelenvelope.ParseProviderReportedCostUSD(row.ProviderReportedCostUsd) + if row.State != "started" { + if row.State != string(state) || valueOrZero(row.RecoveryKind) != string(failure.Recovery) || + row.ApiFormat != args.ApiFormat || + row.ApiVariant != args.ApiVariant || + row.ProviderRequestID != args.RequestID || + row.ProviderResponseID != args.ResponseID || + row.ErrorKind != string(failure.ErrorKind) || + row.ErrorCode != failure.ErrorCode || + row.ErrorMessage != failure.ErrorMessage || !equalJSON(row.ErrorDetails, details) || + storedCost != evidence.Cost || + valueOrZero(row.InputTokensTotal) != valueOrZero(args.InputTokens) || + valueOrZero(row.UncachedInputTokens) != valueOrZero(args.UncachedTokens) || + valueOrZero(row.CacheReadInputTokens) != valueOrZero(args.CacheReadTokens) || + valueOrZero(row.CacheWriteInputTokens) != valueOrZero(args.CacheWriteTokens) || + valueOrZero(row.OutputTokensTotal) != valueOrZero(args.OutputTokens) || + valueOrZero(row.ReasoningOutputTokens) != valueOrZero(args.ReasoningTokens) || + !equalJSON(row.ProviderMetadata, args.ProviderMetadata) { + return false, storeerr.ErrIdempotencyConflict + } + return false, nil + } + if failure.Recovery == RecoveryRetry && (row.AttemptNumber > 8 || failure.RetryDelay < 0) { + return false, storeerr.ErrStateTransitionConflict + } + if failure.Recovery == RecoveryRetry { + delay := failure.RetryDelay.Microseconds() + args.RetryMicroseconds = &delay + } + args.State = string(state) + args.Recovery = optionalText(string(failure.Recovery)) + args.ErrorKind = string(failure.ErrorKind) + args.ErrorCode = failure.ErrorCode + args.ErrorMessage = failure.ErrorMessage + args.ErrorDetails = details + args.AgentID = h.route.AgentID + args.ID = row.ID + args.RuntimeID = failure.RuntimeLockID + n, err := executiondb.New().FinishExecutionAttempt(ctx, h.unit.DB(), args) + if err != nil { + return false, err + } + if n != 1 { + return false, storeerr.ErrRuntimeLockInactive + } + m.changed() + return true, nil +} + +func (h *Handle) FailModel(ctx context.Context, input ModelFailure) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + if input.ErrorKind == "" || input.ErrorMessage == "" || + input.Recovery != RecoveryRetry && input.Recovery != RecoveryNone { + return false, errors.New("model failure requires error and retry or terminal disposition") + } + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return false, err + } + if input.Recovery == RecoveryNone { + return h.terminalFailure(ctx, m, input) + } + row, err := executiondb.New(). + ReadExecutionAttempt(ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, + ID: input.ContextID}) + if err != nil { + return false, err + } + return h.finishAttempt(ctx, m, row, ContextFailed, input, input.Evidence) + }) +} + +func (h *Handle) terminalFailure( + ctx context.Context, + m *executionMutation, + input ModelFailure, +) (bool, error) { + q := executiondb.New() + row, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: input.ContextID}, + ) + if err != nil { + return false, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return false, err + } + changed, err := h.finishAttempt(ctx, m, row, ContextFailed, input, input.Evidence) + if err != nil { + return false, err + } + return h.publishTerminalFailure(ctx, m, input, row, snapshot, changed) +} + +func (h *Handle) publishTerminalFailure(ctx context.Context, m *executionMutation, input ModelFailure, + row executiondb.ReadExecutionAttemptRow, snapshot ExecutionSnapshot, changed bool) (bool, error) { + q := executiondb.New() + parts := []Content{{Kind: "error", Text: input.ErrorMessage}} + if !changed { + output, err := q.FindExecutionOutput( + ctx, + h.unit.DB(), + executiondb.FindExecutionOutputParams{AgentID: h.route.AgentID, ContextID: input.ContextID}, + ) + if err != nil { + return false, err + } + if output.StopReason != "error" || output.ServedProviderModelSlug != input.ServedModel { + return false, storeerr.ErrIdempotencyConflict + } + return false, h.contentMatches(ctx, uuid.Nil, output.ID, parts) + } + output, err := q.CreateExecutionOutput(ctx, h.unit.DB(), executiondb.CreateExecutionOutputParams{ + AgentID: h.route.AgentID, ContextID: input.ContextID, StopReason: "error", ServedModel: input.ServedModel}) + if err != nil { + return false, err + } + event, err := h.appendEvent(ctx, row.TurnID, "model_output", uuid.Nil, output.ID, uuid.Nil, false) + if err != nil { + return false, err + } + if err = h.writeContent(ctx, "model_output", output.ID, parts); err != nil { + return false, err + } + if err = h.advanceTurn(ctx, event, true); err != nil { + return false, err + } + if err = applyAcceptedOutput(m, + snapshot, + attemptRecord(row).Context, + output.ID, + event, + false, + false); err != nil { + return false, err + } + if err = h.parentEffect(ctx, + "failed", + input.ErrorMessage, + "model_output:"+output.ID.String(), + uuid.Nil); err != nil { + return false, err + } + return true, nil +} diff --git a/internal/storage/internal/agentexecution/commands_family_integration_test.go b/internal/storage/internal/agentexecution/commands_family_integration_test.go new file mode 100644 index 000000000..64b7b0698 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_family_integration_test.go @@ -0,0 +1,317 @@ +//go:build integration + +package agentexecution_test + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/interactionform" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/stretchr/testify/require" +) + +func TestExecutionCreatedAgentPublicationAndSavepoint(t *testing.T) { + f := migratedExecutionFixture(t) + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = u.Rollback(context.Background()) }() + h, err := u.CreateAgent( + t.Context(), + agentexecution.CreateAgentInput{ID: uuid.New(), ProjectID: f.project, ConfigID: f.config}, + ) + require.NoError(t, err) + route, err := h.Route() + require.NoError(t, err) + receiveCommand(t, h, "queued", "created") + assertCommandState(t, u, h) + var heads int + require.NoError( + t, + u.DB().QueryRow(t.Context(), `SELECT count(*) FROM agent_execution_state`).Scan(&heads), + ) + require.Zero(t, heads) + sentinel := errors.New("rollback child") + var escaped *agentexecution.Handle + require.ErrorIs(t, u.Savepoint(t.Context(), func(u *agentexecution.Unit) error { + escaped, err = u.CreateAgent( + t.Context(), + agentexecution.CreateAgentInput{ + ID: uuid.New(), + ProjectID: f.project, + ParentID: route.AgentID, + ConfigID: f.config, + SubagentKey: "child", + }, + ) + require.NoError(t, err) + receiveCommand(t, escaped, "queued", "discarded") + return sentinel + }), sentinel) + _, err = escaped.Route() + require.ErrorIs(t, err, agentexecution.ErrUnitClosed) + statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) + require.NoError(t, err) + require.Equal(t, 2, statements) + require.Equal(t, 2, roundTrips) + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_execution_state`).Scan(&heads), + ) + require.Equal(t, 1, heads) +} + +func familyFixture( + t *testing.T, + childFirst bool, +) (executionFixture, *agentexecution.Unit, *agentexecution.Handle, *agentexecution.Handle, uuid.UUID) { + t.Helper() + f := migratedExecutionFixture(t) + parent := f.create(t, uuid.New(), uuid.Nil) + child := uuid.MustParse("ffffffff-ffff-4fff-bfff-ffffffffffff") + if childFirst { + child = uuid.MustParse("00000000-0000-4000-8000-000000000001") + } + f.create(t, child, parent) + lease := uuid.New() + _, err := f.pool.Exec( + t.Context(), + `INSERT INTO agent_runtime_locks(id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + VALUES($1,$2,$3,statement_timestamp(),statement_timestamp(),statement_timestamp()+interval '90 seconds')`, + lease, + child, + uuid.New(), + ) + require.NoError(t, err) + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(t.Context()) + require.NoError(t, err) + t.Cleanup(func() { _ = u.Rollback(context.Background()) }) + require.NoError( + t, + u.LockAgentRefs( + t.Context(), + []lifecyclelock.AgentRef{ + {ProjectID: f.project, AgentID: parent}, + {ProjectID: f.project, AgentID: child}, + }, + agentexecution.LifecycleAuthority{}, + ), + ) + p, err := u.Agent( + agentexecution.AgentRoute{CellID: "test", ProjectID: f.project, AgentID: parent, RootAgentID: parent}, + ) + require.NoError(t, err) + c, err := u.Agent( + agentexecution.AgentRoute{CellID: "test", ProjectID: f.project, AgentID: child, RootAgentID: parent}, + ) + require.NoError(t, err) + return f, u, p, c, lease +} + +func TestExecutionParentEffects(t *testing.T) { + for _, order := range []bool{false, true} { + for _, kind := range []string{"answer", "question", "failure"} { + t.Run( + map[bool]string{false: "parent_first", true: "child_first"}[order]+"/"+kind, + func(t *testing.T) { + _, u, p, c, lease := familyFixture(t, order) + if kind == "question" { + ref := toolCommand(t, c, lease, "built_in") + _, err := c.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + form := interactionform.Form{ + Title: "Parent question", + Questions: []interactionform.Question{ + {Prompt: "Which?", Options: []interactionform.Option{{Label: "yes"}}}, + }, + } + _, err = c.OpenInteraction( + t.Context(), + agentexecution.OpenInteractionInput{ToolRef: ref, Question: &form}, + ) + require.NoError(t, err) + } else { + receiveCommand(t, c, "queued", "input") + _, err := c.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, c, lease) + if kind == "answer" { + input := agentexecution.AcceptOutputInput{RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: commandResponse(modelenvelope.StopReasonEndTurn)} + _, err = c.AcceptOutput(t.Context(), input) + require.NoError(t, err) + replay, err := c.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.False(t, replay.Created) + } else { + _, + err = c.FailModel(t.Context(), + agentexecution.ModelFailure{ContextID: prepared.Context.ID, + RuntimeLockID: lease, + ErrorKind: "runtime", + ErrorMessage: "failure"}) + require.NoError(t, err) + } + } + pr, err := p.Route() + require.NoError(t, err) + parentState := assertCommandState(t, u, p) + require.Equal(t, agentexecution.AdmitAllSteering, parentState.Selection.Admission) + assertCommandState(t, u, c) + var inputs int + require.NoError( + t, + u.DB(). + QueryRow(t.Context(), `SELECT count(*) FROM agent_inputs WHERE agent_id=$1`, pr.AgentID). + Scan(&inputs), + ) + require.Equal(t, 1, inputs) + statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) + require.NoError(t, err) + require.Equal(t, 4, statements) + require.Equal(t, 2, roundTrips) + }, + ) + } + } +} + +func TestExecutionArchiveFamily(t *testing.T) { + f, u, p, c, lease := familyFixture(t, true) + toolCommand(t, c, lease, "custom") + receiveCommand(t, p, "queued", "parent") + receiveCommand(t, c, "queued", "child") + pr, err := p.Route() + require.NoError(t, err) + cr, err := c.Route() + require.NoError(t, err) + require.NoError(t, u.ArchiveAgents(t.Context(), []agentexecution.AgentRoute{pr, cr}, uuid.Nil)) + for _, h := range []*agentexecution.Handle{p, c} { + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.WaitArchived, snapshot.Selection.Wait) + replay, err := h.Archive(t.Context(), uuid.Nil) + require.NoError(t, err) + require.False(t, replay) + } + require.NoError(t, u.Commit(t.Context(), "archive family")) + var wakeups int + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_wakeups`).Scan(&wakeups)) + require.Zero(t, wakeups) +} + +func familyParentRuntime(t *testing.T, u *agentexecution.Unit, p *agentexecution.Handle) uuid.UUID { + t.Helper() + route, err := p.Route() + require.NoError(t, err) + lease := uuid.New() + _, + err = u.DB().Exec(t.Context(), + `INSERT INTO agent_runtime_locks(id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + VALUES($1,$2,$3,statement_timestamp(),statement_timestamp(),statement_timestamp()+interval '90 seconds')`, + lease, route.AgentID, uuid.New()) + require.NoError(t, err) + return lease +} + +func TestExecutionChildCommands(t *testing.T) { + for _, order := range []bool{false, true} { + for _, mode := range []string{"send", "cancel"} { + t.Run( + map[bool]string{false: "parent_first", true: "child_first"}[order]+"/"+mode, + func(t *testing.T) { + _, u, p, c, lease := familyFixture(t, order) + parentLease := familyParentRuntime(t, u, p) + parentRef := toolCommand(t, p, parentLease, "built_in") + _, err := p.AuthorizeTool(t.Context(), parentRef) + require.NoError(t, err) + childRef := toolCommand(t, c, lease, "built_in") + _, err = c.AuthorizeTool(t.Context(), childRef) + require.NoError(t, err) + form := interactionform.Form{Title: "Question", + Questions: []interactionform.Question{{Prompt: "Continue?", + Options: []interactionform.Option{{Label: "yes"}}}}} + question, + err := c.OpenInteraction(t.Context(), + agentexecution.OpenInteractionInput{ToolRef: childRef, + Question: &form}) + require.NoError(t, err) + cr, err := c.Route() + require.NoError(t, err) + completion := agentexecution.ToolCompletion{ToolRef: parentRef, + Outcome: "succeeded", + Content: []agentexecution.Content{{Kind: "text", + Text: "done"}}} + run := func() (agentexecution.CompletedTool, error) { + if mode == "send" { + return p.SendToChild(t.Context(), cr, uuid.Nil, "message", completion) + } + return p.CancelChild(t.Context(), cr, uuid.Nil, completion) + } + result, err := run() + require.NoError(t, err) + require.True(t, result.Changed) + var state string + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT state FROM agent_interactions WHERE id=$1`, + question.ID).Scan(&state)) + require.Equal(t, "canceled", state) + later := receiveCommand(t, c, "steering", "after") + replay, err := run() + require.NoError(t, err) + require.False(t, replay.Changed) + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT state FROM agent_inputs WHERE id=$1`, + later.ID).Scan(&state)) + require.Equal(t, "received", state) + assertCommandState(t, u, p) + assertCommandState(t, u, c) + require.NoError(t, u.Commit(t.Context(), "child command")) + }, + ) + } + } +} + +func TestExecutionParentReportPreservesInteraction(t *testing.T) { + _, u, p, c, lease := familyFixture(t, true) + parentLease := familyParentRuntime(t, u, p) + ref := toolCommand(t, p, parentLease, "built_in") + _, err := p.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + form := interactionform.Form{Title: "Question", + Questions: []interactionform.Question{{Prompt: "Choose", + Options: []interactionform.Option{{Label: "yes"}}}}} + opened, err := p.OpenInteraction( + t.Context(), + agentexecution.OpenInteractionInput{ToolRef: ref, Question: &form}, + ) + require.NoError(t, err) + receiveCommand(t, c, "queued", "answer") + _, err = c.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, c, lease) + _, + err = c.AcceptOutput(t.Context(), + agentexecution.AcceptOutputInput{RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: commandResponse(modelenvelope.StopReasonEndTurn)}) + require.NoError(t, err) + var state string + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT state FROM agent_interactions WHERE id=$1`, + opened.ID).Scan(&state)) + require.Equal(t, "open", state) + snapshot := assertCommandState(t, u, p) + require.Equal(t, agentexecution.WaitToolBatch, snapshot.Selection.Wait) + require.NoError(t, u.Commit(t.Context(), "parent report")) +} diff --git a/internal/storage/internal/agentexecution/commands_inputs.go b/internal/storage/internal/agentexecution/commands_inputs.go new file mode 100644 index 000000000..2b789cb34 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_inputs.go @@ -0,0 +1,325 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "math" + "slices" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ReceiveContentInput struct { + ID uuid.UUID + ActorID uuid.UUID + IntegrationTargetID uuid.UUID + DeliveryMode string + IdempotencyScope string + IdempotencyKey string + Metadata json.RawMessage + Content []Content +} + +type ReceivedContent struct { + ID uuid.UUID + Created bool +} + +type AdmittedContent struct { + ActorID uuid.UUID + ID uuid.UUID + IntegrationTargetID uuid.UUID + Event ExecutionEvent +} + +type InputAdmission struct { + TurnID uuid.UUID + Inputs []AdmittedContent +} + +func (h *Handle) ReceiveContent(ctx context.Context, input ReceiveContentInput) (ReceivedContent, error) { + return executeCommand(ctx, h, func(m *executionMutation) (ReceivedContent, error) { + if input.DeliveryMode != "queued" && input.DeliveryMode != "steering" { + return ReceivedContent{}, errors.New("invalid content delivery mode") + } + if input.IdempotencyKey != "" && input.IdempotencyScope == "" { + return ReceivedContent{}, errors.New("input idempotency key requires a scope") + } + for _, part := range input.Content { + if part.Kind != "text" && part.Kind != "artifact" && part.Kind != "structured_data" { + return ReceivedContent{}, errors.New("invalid input content") + } + } + metadata, err := objectJSON(input.Metadata) + if err != nil { + return ReceivedContent{}, err + } + q := executiondb.New() + existing, err := q.FindExecutionInput( + ctx, + h.unit.DB(), + executiondb.FindExecutionInputParams{AgentID: h.route.AgentID, + ID: input.ID, Scope: optionalText(input.IdempotencyScope), Key: optionalText(input.IdempotencyKey)}, + ) + if err == nil { + if existing.InputKind != "content" || existing.DeliveryMode != input.DeliveryMode || + valueOrZero(existing.ActorID) != input.ActorID || + valueOrZero(existing.IntegrationTargetID) != input.IntegrationTargetID || + !equalJSON(existing.Metadata, metadata) { + return ReceivedContent{}, storeerr.ErrIdempotencyConflict + } + return ReceivedContent{ + ID: existing.ID, + }, h.contentMatches( + ctx, + existing.ID, + uuid.Nil, + input.Content, + ) + } + if !errors.Is(err, pgx.ErrNoRows) { + return ReceivedContent{}, err + } + if input.ID == uuid.Nil { + input.ID = uuid.New() + } + row, err := q.ReceiveExecutionInput(ctx, h.unit.DB(), executiondb.ReceiveExecutionInputParams{ + ID: input.ID, AgentID: h.route.AgentID, ProjectID: h.route.ProjectID, ActorID: nullableID(input.ActorID), + TargetID: nullableID( + input.IntegrationTargetID, + ), Mode: input.DeliveryMode, Kind: "content", Metadata: metadata, + Scope: optionalText(input.IdempotencyScope), Key: optionalText(input.IdempotencyKey)}) + if err != nil { + return ReceivedContent{}, err + } + if err = h.writeContent(ctx, "agent_input", row.ID, input.Content); err != nil { + return ReceivedContent{}, err + } + m.changed() + return ReceivedContent{ID: row.ID, Created: true}, nil + }) +} + +func (h *Handle) AdmitInputs(ctx context.Context) (InputAdmission, error) { + return executeCommand(ctx, h, func(m *executionMutation) (InputAdmission, error) { + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return InputAdmission{}, err + } + if snapshot.Selection.Work != WorkInput { + return InputAdmission{}, nil + } + return h.admitInputs(ctx, m, snapshot.Selection.Admission) + }) +} + +func (h *Handle) admitInputs( + ctx context.Context, + m *executionMutation, + kind AdmissionKind, +) (InputAdmission, error) { + q := executiondb.New() + rows, err := q.ListExecutionBacklog( + ctx, + h.unit.DB(), + executiondb.ListExecutionBacklogParams{AgentID: h.route.AgentID}, + ) + if err != nil { + return InputAdmission{}, err + } + result := InputAdmission{TurnID: uuid.New()} + for _, row := range rows { + if kind == AdmitAllSteering && row.DeliveryMode != "steering" || + kind == AdmitOneQueued && row.DeliveryMode != "queued" { + continue + } + event, err := h.appendEvent(ctx, result.TurnID, "agent_input", row.ID, uuid.Nil, uuid.Nil, true) + if err != nil { + return InputAdmission{}, err + } + n, err := q.ResolveExecutionInput(ctx, h.unit.DB(), executiondb.ResolveExecutionInputParams{ + AgentID: h.route.AgentID, ID: row.ID, EventID: &event.ID}) + if err != nil { + return InputAdmission{}, err + } + if n != 1 { + return InputAdmission{}, storeerr.ErrStateTransitionConflict + } + result.Inputs = append( + result.Inputs, + AdmittedContent{ + ID: row.ID, + ActorID: valueOrZero(row.ActorID), + IntegrationTargetID: valueOrZero(row.IntegrationTargetID), + Event: event, + }, + ) + if kind == AdmitOneQueued { + break + } + } + if len(result.Inputs) == 0 { + return InputAdmission{}, nil + } + last := result.Inputs[len(result.Inputs)-1].Event + if err := q.OpenExecutionTurn(ctx, h.unit.DB(), executiondb.OpenExecutionTurnParams{ + AgentID: h.route.AgentID, ID: result.TurnID, EventID: last.ID}); err != nil { + return InputAdmission{}, err + } + m.openTurn(result.TurnID) + return result, nil +} + +type BacklogChange struct { + ID uuid.UUID + DeliveryMode string + Cancel bool + BeforeID uuid.UUID + Position string +} + +func (h *Handle) ChangeBacklog(ctx context.Context, change BacklogChange) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + if change.ID == uuid.Nil || change.DeliveryMode != "queued" && change.DeliveryMode != "steering" { + return false, errors.New("invalid backlog change") + } + switch change.Position { + case "", "front", "back": + if change.BeforeID != uuid.Nil { + return false, errors.New("backlog change does not accept an anchor") + } + case "before", "after": + if change.BeforeID == uuid.Nil || change.BeforeID == change.ID { + return false, storeerr.ErrStateTransitionConflict + } + default: + return false, errors.New("invalid backlog position") + } + if change.Cancel && change.Position != "" || + (change.Cancel || change.Position != "") && change.DeliveryMode != "queued" { + return false, errors.New("only queued backlog supports cancellation and movement") + } + q := executiondb.New() + rows, err := q.ListExecutionBacklog(ctx, h.unit.DB(), + executiondb.ListExecutionBacklogParams{AgentID: h.route.AgentID}) + if err != nil { + return false, err + } + index := slices.IndexFunc( + rows, + func(row executiondb.ListExecutionBacklogRow) bool { return row.ID == change.ID }, + ) + if index < 0 { + if change.DeliveryMode == "steering" { + row, err := q.FindExecutionInput(ctx, h.unit.DB(), + executiondb.FindExecutionInputParams{AgentID: h.route.AgentID, ID: change.ID}) + if err == nil && row.InputKind == "content" && row.State == "resolved" && + row.AdmittedEventID != nil { + return false, nil + } + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return false, err + } + } + return false, storeerr.ErrStateTransitionConflict + } + selected := rows[index] + if change.Cancel { + if selected.DeliveryMode != "queued" { + return false, storeerr.ErrStateTransitionConflict + } + return h.changeBacklogRow(ctx, m, selected, true) + } + if change.Position != "" && selected.DeliveryMode != "queued" { + return false, storeerr.ErrStateTransitionConflict + } + if change.Position == "" && change.DeliveryMode == selected.DeliveryMode { + if change.DeliveryMode == "steering" { + return false, nil + } + return false, storeerr.ErrStateTransitionConflict + } + rows = slices.DeleteFunc(rows, func(row executiondb.ListExecutionBacklogRow) bool { + return row.DeliveryMode != change.DeliveryMode + }) + rank, available, err := backlogRank(rows, change) + if err != nil { + return false, err + } + var changed bool + if !available { + for i := range rows { + rows[i].InputRank = int64(i+1) * 1024 + updated, err := h.changeBacklogRow(ctx, m, rows[i], false) + if err != nil { + return false, err + } + changed = changed || updated + } + rank, available, err = backlogRank(rows, change) + if err != nil { + return false, err + } + if !available { + return false, storeerr.ErrStateTransitionConflict + } + } + selected.DeliveryMode = change.DeliveryMode + selected.InputRank = rank + updated, err := h.changeBacklogRow(ctx, m, selected, false) + return changed || updated, err + }) +} + +func (h *Handle) changeBacklogRow(ctx context.Context, m *executionMutation, + row executiondb.ListExecutionBacklogRow, cancel bool) (bool, error) { + n, err := executiondb.New().ChangeExecutionInput(ctx, h.unit.DB(), executiondb.ChangeExecutionInputParams{ + AgentID: h.route.AgentID, ID: row.ID, Mode: row.DeliveryMode, Rank: row.InputRank, Cancel: cancel}) + if n > 0 { + m.changed() + } + return n > 0, err +} + +func backlogRank(rows []executiondb.ListExecutionBacklogRow, change BacklogChange) (int64, bool, error) { + neighbors := make([]int64, 0, len(rows)) + anchor := -1 + for _, row := range rows { + if row.ID == change.ID { + continue + } + if row.ID == change.BeforeID { + anchor = len(neighbors) + } + neighbors = append(neighbors, row.InputRank) + } + destination := len(neighbors) + switch change.Position { + case "front": + destination = 0 + case "before", "after": + if anchor < 0 { + return 0, false, storeerr.ErrStateTransitionConflict + } + destination = anchor + if change.Position == "after" { + destination++ + } + } + var lower int64 + if destination > 0 { + lower = neighbors[destination-1] + } + if destination < len(neighbors) { + upper := neighbors[destination] + return lower + (upper-lower)/2, upper-lower > 1, nil + } + if lower > math.MaxInt64-1024 { + return 0, false, errors.New("backlog rank overflow") + } + return lower + 1024, true, nil +} diff --git a/internal/storage/internal/agentexecution/commands_integration_test.go b/internal/storage/internal/agentexecution/commands_integration_test.go new file mode 100644 index 000000000..c066538ac --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_integration_test.go @@ -0,0 +1,959 @@ +//go:build integration + +package agentexecution_test + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +func commandFixture(t *testing.T) (executionFixture, uuid.UUID, uuid.UUID) { + t.Helper() + f := migratedExecutionFixture(t) + a := f.create(t, uuid.New(), uuid.Nil) + lease := uuid.New() + _, err := f.pool.Exec( + t.Context(), + `INSERT INTO agent_runtime_locks(id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + VALUES($1,$2,$3,statement_timestamp(),statement_timestamp(),statement_timestamp()+interval '90 seconds')`, + lease, + a, + uuid.New(), + ) + require.NoError(t, err) + return f, a, lease +} + +func assertCommandState( + t *testing.T, + u *agentexecution.Unit, + h *agentexecution.Handle, +) agentexecution.ExecutionSnapshot { + t.Helper() + got, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + rebuilt, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + left, right := got.Head, rebuilt.Head + if got.Selection.Work == agentexecution.WorkInput || got.Selection.Work == agentexecution.WorkTool { + left.LogicalReadyAt = nil + right.LogicalReadyAt = nil + } + require.Equal(t, left, right) + + _, err = u.DB().Exec(t.Context(), "SET CONSTRAINTS ALL IMMEDIATE") + require.NoError(t, err) + _, err = u.DB().Exec(t.Context(), "SET CONSTRAINTS ALL DEFERRED") + require.NoError(t, err) + return got +} + +func receiveCommand( + t *testing.T, + h *agentexecution.Handle, + mode, text string, +) agentexecution.ReceivedContent { + t.Helper() + row, err := h.ReceiveContent( + t.Context(), + agentexecution.ReceiveContentInput{ID: uuid.New(), DeliveryMode: mode, + Content: []agentexecution.Content{{Kind: "text", Text: text}}}, + ) + require.NoError(t, err) + require.True(t, row.Created) + return row +} + +func prepareCommand(t *testing.T, h *agentexecution.Handle, lease uuid.UUID) agentexecution.PreparedModel { + t.Helper() + snapshot, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + require.NotNil(t, snapshot.Selection.Model) + prepared, err := h.PrepareModel( + t.Context(), + agentexecution.PrepareModelInput{RuntimeLockID: lease, Selected: *snapshot.Selection.Model}, + ) + require.NoError(t, err) + require.True(t, prepared.Claimed) + return prepared +} + +func commandResponse(reason modelenvelope.StopReason) modelenvelope.ResponseEnvelope { + return modelenvelope.ResponseEnvelope{ + RequestedProviderModelSlug: "test", + APIFormat: "openai-responses", + APIVariant: "default", + Normalized: modelenvelope.ResponseNormalized{ + ID: "response", + StopReason: reason, + Content: []modelenvelope.ResponsePart{{Type: "text", Text: "answer"}}, + }, + } +} + +func TestExecutionSemanticCommands(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + config, err := h.ActivateConfig( + t.Context(), + agentexecution.ActivateConfigInput{ID: uuid.New(), ConfigID: f.config, IdempotencyKey: "initial"}, + ) + require.NoError(t, err) + require.True(t, config.Created) + assertCommandState(t, u, h) + first := receiveCommand(t, h, "queued", "first") + second := receiveCommand(t, h, "queued", "second") + assertCommandState(t, u, h) + changed, err := h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ID: second.ID, DeliveryMode: "steering"}, + ) + require.NoError(t, err) + require.True(t, changed) + assertCommandState(t, u, h) + changed, err = h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ID: first.ID, DeliveryMode: "queued", Cancel: true}, + ) + require.NoError(t, err) + require.True(t, changed) + assertCommandState(t, u, h) + admitted, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + require.Len(t, admitted.Inputs, 1) + require.Equal(t, second.ID, admitted.Inputs[0].ID) + snapshot := assertCommandState(t, u, h) + selected := *snapshot.Selection.Model + prepared := prepareCommand(t, h, lease) + require.Equal(t, []uuid.UUID{second.ID}, prepared.Context.Opening.InputIDs) + assertCommandState(t, u, h) + replay, err := h.PrepareModel( + t.Context(), + agentexecution.PrepareModelInput{RuntimeLockID: lease, Selected: selected}, + ) + require.NoError(t, err) + require.False(t, replay.Created) + require.Equal(t, prepared.Context.ID, replay.Context.ID) + failure := agentexecution.ModelFailure{ + ContextID: prepared.Context.ID, + RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, + ErrorKind: "transient", + ErrorMessage: "retry", + RetryDelay: time.Second, + } + changed, err = h.FailModel(t.Context(), failure) + require.NoError(t, err) + require.True(t, changed) + waiting := assertCommandState(t, u, h) + require.Equal(t, agentexecution.WaitModelDeadline, waiting.Selection.Wait) + changed, err = h.FailModel(t.Context(), failure) + require.NoError(t, err) + require.False(t, changed) + require.NoError(t, u.Commit(t.Context(), "semantic retry")) + u, h = f.handle(t, a) + require.Eventually(t, func() bool { + snapshot, err := h.LoadExecution(t.Context()) + return err == nil && snapshot.Selection.Work == agentexecution.WorkModel + }, 2*time.Second, 10*time.Millisecond) + retried := prepareCommand(t, h, lease) + require.Equal(t, int32(2), retried.Context.Attempt) + require.Equal(t, prepared.Context.Opening, retried.Context.Opening) + assertCommandState(t, u, h) + input := agentexecution.AcceptOutputInput{ + ContextID: retried.Context.ID, + RuntimeLockID: lease, + Response: commandResponse(modelenvelope.StopReasonMaxTokens), + } + beforeOutput := u.Notifications().Clone() + accepted, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.True(t, accepted.Created) + assertCommandState(t, u, h) + afterOutput := u.Notifications().Clone() + require.NotEqual(t, beforeOutput, afterOutput) + replayed, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.False(t, replayed.Created) + require.Equal(t, accepted.ID, replayed.ID) + require.Equal(t, afterOutput, u.Notifications().Clone()) + continuation := prepareCommand(t, h, lease) + assertCommandState(t, u, h) + _, err = h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ContextID: continuation.Context.ID, RuntimeLockID: lease, + Response: commandResponse(modelenvelope.StopReasonEndTurn)}, + ) + require.NoError(t, err) + final := assertCommandState(t, u, h) + require.Equal(t, agentexecution.WorkNone, final.Selection.Work) + require.NoError(t, u.Commit(t.Context(), "semantic completion")) +} + +func TestExecutionCommandSavepointAndReplay(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + input := agentexecution.ReceiveContentInput{ + ID: uuid.New(), + DeliveryMode: "queued", + IdempotencyScope: "test", + IdempotencyKey: "one", + Content: []agentexecution.Content{{Kind: "text", Text: "one"}}, + } + first, err := h.ReceiveContent(t.Context(), input) + require.NoError(t, err) + expected := assertCommandState(t, u, h) + sentinel := errors.New("rollback") + require.ErrorIs(t, u.Savepoint(t.Context(), func(u *agentexecution.Unit) error { + receiveCommand(t, h, "steering", "later") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + assertCommandState(t, u, h) + return sentinel + }), sentinel) + got := assertCommandState(t, u, h) + require.Equal(t, expected.Head.CurrentTurnID, got.Head.CurrentTurnID) + require.Equal(t, agentexecution.AdmitOneQueued, got.Selection.Admission) + replay, err := h.ReceiveContent(t.Context(), input) + require.NoError(t, err) + require.False(t, replay.Created) + require.Equal(t, first.ID, replay.ID) + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + input.Content[0].Text = "different" + _, err := h.ReceiveContent(t.Context(), input) + return err + }), storeerr.ErrIdempotencyConflict) + require.NoError(t, u.Commit(t.Context(), "savepoint replay")) + u, h = f.handle(t, a) + input.Content[0].Text = "one" + replay, err = h.ReceiveContent(t.Context(), input) + require.NoError(t, err) + require.False(t, replay.Created) + require.NoError(t, u.Commit(t.Context(), "unchanged replay")) +} + +func TestExecutionCompactionCommands(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + _, err := h.ActivateConfig(t.Context(), agentexecution.ActivateConfigInput{ConfigID: f.config}) + require.NoError(t, err) + receiveCommand(t, h, "queued", "input") + _, err = h.AdmitInputs(t.Context()) + require.NoError(t, err) + normal := prepareCommand(t, h, lease) + start, err := h.BeginCompaction( + t.Context(), + agentexecution.BeginCompactionInput{SourceEnd: normal.Context.InputEventSequence, + Failure: agentexecution.ModelFailure{ + ContextID: normal.Context.ID, + RuntimeLockID: lease, + ErrorKind: "context_length", + ErrorMessage: "compact", + }}, + ) + require.NoError(t, err) + require.True(t, start.Model.Created) + assertCommandState(t, u, h) + _, err = h.FailModel(t.Context(), agentexecution.ModelFailure{ContextID: start.Model.Context.ID, + RuntimeLockID: lease, Recovery: agentexecution.RecoveryRetry, ErrorKind: "transient", ErrorMessage: "retry"}) + require.NoError(t, err) + retry := prepareCommand(t, h, lease) + snapshot := assertCommandState(t, u, h) + require.Equal(t, retry.Context.ID, snapshot.Head.CompactionContextID) + require.Equal(t, start.Model.Context.Opening, retry.Context.Opening) + reduced, err := h.BeginCompaction(t.Context(), agentexecution.BeginCompactionInput{SourceEnd: 1, + Failure: agentexecution.ModelFailure{ + ContextID: retry.Context.ID, + RuntimeLockID: lease, + ErrorKind: "context_length", + ErrorMessage: "reduce", + }}) + require.NoError(t, err) + require.True(t, reduced.Model.Created) + assertCommandState(t, u, h) + cp := agentexecution.PublishCheckpointInput{ + ContextID: reduced.Model.Context.ID, + RuntimeLockID: lease, + Summary: "summary", + Evidence: agentexecution.ModelEvidence{APIFormat: "openai-responses", APIVariant: "default"}, + } + checkpoint, err := h.PublishCheckpoint(t.Context(), cp) + require.NoError(t, err) + require.True(t, checkpoint.Created) + snapshot = assertCommandState(t, u, h) + require.Equal(t, checkpoint.ID, snapshot.Head.PendingCheckpointID) + again, err := h.PublishCheckpoint(t.Context(), cp) + require.NoError(t, err) + require.False(t, again.Created) + next := prepareCommand(t, h, lease) + require.Equal(t, normal.Context.Opening, next.Context.Opening) + assertCommandState(t, u, h) + changed, err := h.FailModel( + t.Context(), + agentexecution.ModelFailure{ContextID: next.Context.ID, RuntimeLockID: lease, + ErrorKind: "runtime", ErrorMessage: "terminal"}, + ) + require.NoError(t, err) + require.True(t, changed) + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "compaction trace")) +} + +func TestExecutionOutputToolProposalCommand(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = []modelenvelope.ResponsePart{ + {Type: "tool_call", ProviderCallID: "one", ToolName: "list_agents", ToolInput: json.RawMessage(`{}`)}, + { + Type: "tool_call", + ProviderCallID: "two", + ToolName: "list_agents", + ToolInput: json.RawMessage(`{}`), + ToolCallError: "bad arguments", + }, + } + input := agentexecution.AcceptOutputInput{ + ContextID: prepared.Context.ID, + RuntimeLockID: lease, + Response: response, + Tools: []agentexecution.ToolProposal{ + {ProviderCallID: "one", Type: "built_in"}, + {ProviderCallID: "two", Type: "built_in"}, + }, + } + output, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.Len(t, output.Tools, 2) + snapshot := assertCommandState(t, u, h) + require.True(t, snapshot.Selection.IncompleteTools) + again, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.False(t, again.Created) + require.NoError(t, u.Commit(t.Context(), "tool proposal")) +} + +func TestExecutionCommandFencing(t *testing.T) { + for _, mode := range []string{"expired", "canceled", "wrong"} { + t.Run(mode, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + selected, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + if mode == "expired" { + _, err = u.DB(). + Exec(t.Context(), ` +UPDATE agent_runtime_locks SET started_at=statement_timestamp()-interval '3 +minutes',renewed_at=statement_timestamp()-interval '2 +minutes',lease_expires_at=statement_timestamp()-interval '1 minute' WHERE id=$1`, lease) + } else if mode == "canceled" { + _, err = u.DB().Exec(t.Context(), ` +UPDATE agent_runtime_locks SET cancel_requested_at=statement_timestamp() WHERE id=$1`, lease) + } else { + lease = uuid.New() + } + require.NoError(t, err) + _, err = h.PrepareModel( + t.Context(), + agentexecution.PrepareModelInput{RuntimeLockID: lease, Selected: *selected.Selection.Model}, + ) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) + require.ErrorIs(t, u.Commit(context.Background(), "failed fence"), agentexecution.ErrUnitAborted) + }) + } +} + +func TestExecutionPreparationPreservesSourceLineage(t *testing.T) { + for _, tools := range []bool{false, true} { + t.Run(map[bool]string{false: "output_limit", true: "tools"}[tools], func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + older := receiveCommand(t, h, "queued", "older") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + newer := receiveCommand(t, h, "steering", "newer") + _, err = h.AdmitInputs(t.Context()) + require.NoError(t, err) + source := prepareCommand(t, h, lease) + require.Equal(t, []uuid.UUID{older.ID, newer.ID}, source.Context.Opening.InputIDs) + input := agentexecution.AcceptOutputInput{ContextID: source.Context.ID, RuntimeLockID: lease, + Response: commandResponse(modelenvelope.StopReasonMaxTokens)} + if tools { + input.Response.Normalized.StopReason = modelenvelope.StopReasonToolUse + input.Response.Normalized.Content = []modelenvelope.ResponsePart{{Type: "tool_call", + ProviderCallID: "one", ToolName: "list_agents", ToolInput: json.RawMessage(`{}`)}} + input.Tools = []agentexecution.ToolProposal{{ProviderCallID: "one", Type: "built_in"}} + } + output, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + if tools { + ref := agentexecution.ToolRef{ID: output.Tools[0].ID, RuntimeLockID: lease} + _, err = h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + _, err = h.CompleteTool(t.Context(), agentexecution.ToolCompletion{ToolRef: ref, + Outcome: "succeeded", Content: []agentexecution.Content{{Kind: "text", Text: "done"}}}) + require.NoError(t, err) + } + snapshot := assertCommandState(t, u, h) + watermark := source.Context.InputEventSequence + 1 + if tools { + watermark = snapshot.View.ToolBatch.Completion.LastResultSequence + } + require.Greater(t, watermark, source.Context.InputEventSequence) + continuation := prepareCommand(t, h, lease) + require.Equal(t, []uuid.UUID{older.ID, newer.ID}, continuation.Context.Opening.InputIDs) + require.Equal(t, watermark, continuation.Context.InputEventSequence) + _, err = h.FailModel( + t.Context(), + agentexecution.ModelFailure{ContextID: continuation.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, ErrorKind: "transient", ErrorMessage: "retry"}, + ) + require.NoError(t, err) + retry := prepareCommand(t, h, lease) + require.Equal(t, continuation.Context.Opening, retry.Context.Opening) + _, err = h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ContextID: retry.Context.ID, RuntimeLockID: lease, + Response: commandResponse(modelenvelope.StopReasonMaxTokens)}, + ) + require.NoError(t, err) + next := prepareCommand(t, h, lease) + require.Equal(t, continuation.Context.Opening, next.Context.Opening) + rebuilt, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + got, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + require.Equal(t, got.Head, rebuilt.Head) + require.NoError(t, u.Commit(t.Context(), "corrected lineage")) + }) + } +} + +func TestExecutionConfigAndCheckpointBoundaries(t *testing.T) { + for _, boundary := range []string{"config", "steering", "checkpoint"} { + t.Run(boundary, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + if boundary == "config" { + prior := prepareCommand(t, h, lease) + _, err = h.FailModel( + t.Context(), + agentexecution.ModelFailure{ContextID: prior.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, ErrorKind: "transient", ErrorMessage: "retry"}, + ) + require.NoError(t, err) + selected, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + cfg := agentexecution.ActivateConfigInput{ConfigID: f.config, IdempotencyKey: "change"} + _, err = h.ActivateConfig(t.Context(), cfg) + require.NoError(t, err) + assertCommandState(t, u, h) + replay, err := h.ActivateConfig(t.Context(), cfg) + require.NoError(t, err) + require.False(t, replay.Created) + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.PrepareModel( + t.Context(), + agentexecution.PrepareModelInput{ + RuntimeLockID: lease, + Selected: *selected.Selection.Model, + }, + ) + return err + }), storeerr.ErrAgentNotAdvanceable) + } + normal := prepareCommand(t, h, lease) + if boundary == "steering" { + receiveCommand(t, h, "steering", "new boundary") + } + compact, err := h.BeginCompaction( + t.Context(), + agentexecution.BeginCompactionInput{SourceEnd: normal.Context.InputEventSequence, + Failure: agentexecution.ModelFailure{ + ContextID: normal.Context.ID, + RuntimeLockID: lease, + ErrorKind: "context_length", + ErrorMessage: "compact", + }}, + ) + require.NoError(t, err) + if boundary == "steering" { + require.True(t, compact.Preempted) + require.Len(t, compact.Admission.Inputs, 1) + } else { + receiveCommand(t, h, "queued", "later input") + checkpoint, + err := h.PublishCheckpoint(t.Context(), + agentexecution.PublishCheckpointInput{ContextID: compact.Model.Context.ID, + RuntimeLockID: lease, + Summary: "summary", + Evidence: agentexecution.ModelEvidence{APIFormat: "openai-responses", + APIVariant: "default"}}) + require.NoError(t, err) + require.True(t, checkpoint.Created) + } + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "boundary")) + }) + } +} + +func TestExecutionCoalescedPublication(t *testing.T) { + f, a, _ := commandFixture(t) + _, err := f.pool.Exec(t.Context(), `CREATE TABLE execution_publications(kind text,agent_id uuid); + CREATE FUNCTION audit_execution_publication() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN INSERT INTO execution_publications VALUES(TG_TABLE_NAME,NEW.agent_id); RETURN NEW; END; $$; + CREATE TRIGGER audit_head AFTER UPDATE ON agent_execution_state FOR EACH ROW EXECUTE FUNCTION +audit_execution_publication(); + CREATE TRIGGER audit_queue AFTER INSERT OR UPDATE ON agent_wakeups FOR EACH ROW EXECUTE FUNCTION +audit_execution_publication();`) + require.NoError(t, err) + u, h := f.handle(t, a) + _, err = u.DB().Exec(t.Context(), `DELETE FROM agent_runtime_locks WHERE agent_id=$1`, a) + require.NoError(t, err) + first := receiveCommand(t, h, "queued", "first") + second := receiveCommand(t, h, "queued", "second") + _, err = h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ + ID: second.ID, + DeliveryMode: "queued", + Position: "before", + BeforeID: first.ID, + }, + ) + require.NoError(t, err) + var count int + require.NoError( + t, + u.DB().QueryRow(t.Context(), `SELECT count(*) FROM execution_publications`).Scan(&count), + ) + require.Zero(t, count) + require.NoError(t, u.Commit(t.Context(), "coalesced")) + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT count(*) FROM execution_publications`).Scan(&count), + ) + require.Equal(t, 2, count) + var age time.Time + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, a).Scan(&age), + ) + u, h = f.handle(t, a) + third := receiveCommand(t, h, "steering", "third") + require.NoError(t, u.Commit(t.Context(), "preserve ready age")) + var after time.Time + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, a).Scan(&after), + ) + require.True(t, age.Equal(after)) + u, h = f.handle(t, a) + changed, err := h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "steering"}, + ) + require.NoError(t, err) + require.False(t, changed) + var before int + require.NoError( + t, + u.DB().QueryRow(t.Context(), `SELECT count(*) FROM execution_publications`).Scan(&before), + ) + require.NoError(t, u.Commit(t.Context(), "no-op publication")) + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT count(*) FROM execution_publications`).Scan(&count), + ) + require.Equal(t, before, count) +} + +func TestExecutionCheckpointRejectsCutBatch(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + normal := prepareCommand(t, h, lease) + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = []modelenvelope.ResponsePart{ + { + Type: "tool_call", + ProviderCallID: "call", + ToolName: "list_agents", + ToolInput: json.RawMessage(`{}`), + }, + } + output, err := h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ + ContextID: normal.Context.ID, + RuntimeLockID: lease, + Response: response, + Tools: []agentexecution.ToolProposal{{ProviderCallID: "call", Type: "custom"}}, + }, + ) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "tool proposal before checkpoint")) + f.complete(t, a, normal.Context.TurnID, output.Tools[0].ID, 3) + _, err = f.pool.Exec(t.Context(), `UPDATE agents SET next_event_sequence=4 WHERE id=$1`, a) + require.NoError(t, err) + u, h = f.handle(t, a) + next := prepareCommand(t, h, lease) + compact, err := h.BeginCompaction( + t.Context(), + agentexecution.BeginCompactionInput{SourceEnd: 2, Failure: agentexecution.ModelFailure{ + ContextID: next.Context.ID, RuntimeLockID: lease, ErrorKind: "context_length", ErrorMessage: "compact"}}, + ) + require.NoError(t, err) + err = u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.PublishCheckpoint( + t.Context(), + agentexecution.PublishCheckpointInput{ContextID: compact.Model.Context.ID, + RuntimeLockID: lease, + Summary: "unsafe", + Evidence: agentexecution.ModelEvidence{APIFormat: "openai-responses", + APIVariant: "default"}}, + ) + return err + }) + require.ErrorContains(t, err, "boundary is unsafe") + assertCommandState(t, u, h) + var count int + require.NoError( + t, + u.DB(). + QueryRow(t.Context(), `SELECT count(*) FROM context_checkpoints WHERE agent_id=$1`, a). + Scan(&count), + ) + require.Zero(t, count) +} + +func TestExecutionQueueDeadlineAndOwnership(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + require.NoError(t, u.Commit(t.Context(), "owned queue")) + var count int + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_wakeups WHERE agent_id=$1`, a).Scan(&count), + ) + require.Zero(t, count) + u, h = f.handle(t, a) + _, err := u.DB(). + Exec(t.Context(), ` +INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) VALUES($1,statement_timestamp()-interval '1 +hour',statement_timestamp())`, a) + require.NoError(t, err) + _, err = h.AdmitInputs(t.Context()) + require.NoError(t, err) + normal := prepareCommand(t, h, lease) + _, err = h.FailModel( + t.Context(), + agentexecution.ModelFailure{ContextID: normal.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, + RetryDelay: time.Hour, + ErrorKind: "transient", + ErrorMessage: "retry"}, + ) + require.NoError(t, err) + pending, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + _, err = u.DB().Exec(t.Context(), `DELETE FROM agent_runtime_locks WHERE agent_id=$1`, a) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "future delivery")) + var ready time.Time + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, a).Scan(&ready), + ) + require.True(t, ready.Equal(pending.Selection.Model.ReadyAt)) + u, h = f.handle(t, a) + receiveCommand(t, h, "queued", "queued later") + require.NoError(t, u.Commit(t.Context(), "retain retry deadline")) + var later time.Time + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, a).Scan(&later), + ) + require.True(t, ready.Equal(later)) + u, h = f.handle(t, a) + receiveCommand(t, h, "steering", "interrupt") + require.NoError(t, u.Commit(t.Context(), "steering delivery")) + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, a).Scan(&later), + ) + require.True(t, later.Before(ready)) +} + +func TestExecutionPreparationRevalidatesAfterBlockedConfig(t *testing.T) { + for _, boundary := range []string{"retry", "checkpoint"} { + t.Run(boundary, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + model := prepareCommand(t, h, lease) + failure := agentexecution.ModelFailure{ContextID: model.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, ErrorKind: "transient", ErrorMessage: "retry"} + if boundary == "retry" { + _, err = h.FailModel(t.Context(), failure) + require.NoError(t, err) + } else { + compact, + err := h.BeginCompaction(t.Context(), + agentexecution.BeginCompactionInput{Failure: failure, + SourceEnd: model.Context.InputEventSequence}) + require.NoError(t, err) + _, + err = h.PublishCheckpoint(t.Context(), + agentexecution.PublishCheckpointInput{ContextID: compact.Model.Context.ID, + RuntimeLockID: lease, + Summary: "checkpoint", + Evidence: agentexecution.ModelEvidence{APIFormat: "openai-responses", + APIVariant: "default"}}) + require.NoError(t, err) + } + selected, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + require.NotNil(t, selected.Selection.Model) + require.NoError(t, u.Commit(t.Context(), "candidate")) + writer, writerHandle := f.handle(t, a) + _, err = writerHandle.ActivateConfig( + t.Context(), + agentexecution.ActivateConfigInput{ConfigID: f.config}, + ) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + result := make(chan error, 1) + go func() { result <- prepareBlockedCandidate(ctx, f, a, lease, *selected.Selection.Model) }() + require.Eventually(t, func() bool { + var blocked bool + err := f.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_stat_activity + WHERE datname=current_database() AND cardinality(pg_blocking_pids(pid))>0)`).Scan(&blocked) + return err == nil && blocked + }, 5*time.Second, 10*time.Millisecond) + require.NoError(t, writer.Commit(ctx, "concurrent config")) + require.ErrorIs(t, <-result, storeerr.ErrAgentNotAdvanceable) + check, checkHandle := f.handle(t, a) + assertCommandState(t, check, checkHandle) + }) + } +} + +func prepareBlockedCandidate(ctx context.Context, + f executionFixture, + a, + lease uuid.UUID, + decision agentexecution.ModelDecision) error { + cell := agentexecution.NewCell("test", f.pool, nil, nil) + unit, err := cell.Begin(ctx) + if err != nil { + return err + } + defer func() { _ = unit.Rollback(context.WithoutCancel(ctx)) }() + plan, err := unit.PlanAgents(ctx, []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: a}}, + agentexecution.IngressAuthority{}) + if err != nil { + return err + } + if err = unit.LockAgents(ctx, plan); err != nil { + return err + } + handle, + err := unit.Agent(agentexecution.AgentRoute{CellID: cell.ID(), + ProjectID: f.project, + AgentID: a, + RootAgentID: a}) + if err != nil { + return err + } + _, err = handle.PrepareModel( + ctx, + agentexecution.PrepareModelInput{RuntimeLockID: lease, Selected: decision}, + ) + return err +} + +func TestExecutionFenceSamplesTimeAfterRuntimeLock(t *testing.T) { + f, a, lease := commandFixture(t) + unit, h := f.handle(t, a) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + selected, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + require.NoError(t, unit.Commit(t.Context(), "fence candidate")) + _, err = f.pool.Exec(t.Context(), `UPDATE agent_runtime_locks SET + lease_expires_at=statement_timestamp()+interval '2 seconds' WHERE id=$1`, lease) + require.NoError(t, err) + locker, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = locker.Rollback(context.Background()) }() + _, err = locker.Exec(t.Context(), `SELECT id FROM agent_runtime_locks WHERE id=$1 FOR UPDATE`, lease) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + result := make(chan error, 1) + go func() { result <- prepareBlockedCandidate(ctx, f, a, lease, *selected.Selection.Model) }() + require.Eventually(t, func() bool { + var blocked bool + err := f.pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM pg_stat_activity + WHERE datname=current_database() AND cardinality(pg_blocking_pids(pid))>0)`).Scan(&blocked) + return err == nil && blocked + }, time.Second, 10*time.Millisecond) + require.Eventually(t, func() bool { + var expired bool + err := f.pool.QueryRow(ctx, ` +SELECT lease_expires_at<=statement_timestamp() FROM agent_runtime_locks WHERE id=$1`, lease).Scan(&expired) + return err == nil && expired + }, 3*time.Second, 10*time.Millisecond) + require.NoError(t, locker.Commit(ctx)) + require.ErrorIs(t, <-result, storeerr.ErrRuntimeLockInactive) + var contexts int + require.NoError(t, f.pool.QueryRow(ctx, ` +SELECT count(*) FROM model_call_contexts WHERE agent_id=$1`, a).Scan(&contexts)) + require.Zero(t, contexts) +} + +func TestExecutionOutputReplayRejectsConflictingEvidence(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "work") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + input := agentexecution.AcceptOutputInput{ + RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: commandResponse(modelenvelope.StopReasonEndTurn), + } + input.Response.Normalized.Usage.OutputTokens = 10 + accepted, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.True(t, accepted.Created) + for _, variant := range []string{"usage", "metadata", "content"} { + t.Run(variant, func(t *testing.T) { + err := u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + conflict := input + switch variant { + case "usage": + conflict.Response.Normalized.Usage.OutputTokens++ + case "metadata": + conflict.Response.ProviderMetadata.OpenRouter.Provider = "different" + case "content": + conflict.Response.Normalized.Content = []modelenvelope.ResponsePart{ + {Type: "text", Text: "different"}, + } + } + _, err := h.AcceptOutput(t.Context(), conflict) + return err + }) + require.ErrorIs(t, err, storeerr.ErrIdempotencyConflict) + }) + } + replay, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + require.False(t, replay.Created) + require.Equal(t, accepted.ID, replay.ID) + assertCommandState(t, u, h) +} + +func TestExecutionCompactionReplayDoesNotAdmitNewInputs(t *testing.T) { + for _, phase := range []string{"started", "checkpoint", "preempted"} { + t.Run(phase, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "work") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + normal := prepareCommand(t, h, lease) + if phase == "preempted" { + receiveCommand(t, h, "steering", "interrupt") + } + input := agentexecution.BeginCompactionInput{SourceEnd: normal.Context.InputEventSequence, + Failure: agentexecution.ModelFailure{ContextID: normal.Context.ID, RuntimeLockID: lease, + ErrorKind: "context_length", ErrorMessage: "compact"}} + first, err := h.BeginCompaction(t.Context(), input) + require.NoError(t, err) + if phase == "checkpoint" { + _, err = h.PublishCheckpoint(t.Context(), agentexecution.PublishCheckpointInput{ + ContextID: first.Model.Context.ID, RuntimeLockID: lease, Summary: "summary", + Evidence: agentexecution.ModelEvidence{ + APIFormat: "openai-responses", + APIVariant: "default", + }}) + require.NoError(t, err) + } + receiveCommand(t, h, "steering", "later") + before := assertCommandState(t, u, h) + backlog := readCommandBacklog(t, u, a) + replay, err := h.BeginCompaction(t.Context(), input) + require.NoError(t, err) + require.Equal(t, first.Preempted, replay.Preempted) + require.Equal(t, first.Model.Context.ID, replay.Model.Context.ID) + require.False(t, replay.Model.Created) + require.Empty(t, replay.Admission.Inputs) + require.Equal(t, backlog, readCommandBacklog(t, u, a)) + after := assertCommandState(t, u, h) + require.Equal(t, before.Head.CurrentTurnID, after.Head.CurrentTurnID) + }) + } +} + +func TestExecutionContentRejectsDatabaseUnsafeValues(t *testing.T) { + for _, content := range []agentexecution.Content{ + {Kind: "text", Text: "before\x00after"}, + {Kind: "structured_data", Data: json.RawMessage(`{"value":"before\u0000after"}`)}, + {Kind: "text", Text: "safe", Metadata: json.RawMessage(`{"key":"before\u0000after"}`)}, + } { + t.Run(content.Kind+string(content.Metadata), func(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + _, err := h.ReceiveContent( + t.Context(), + agentexecution.ReceiveContentInput{ + ID: uuid.New(), + DeliveryMode: "queued", + Content: []agentexecution.Content{content}, + }, + ) + require.ErrorContains(t, err, "U+0000") + require.Error(t, u.Commit(t.Context(), "unsafe content")) + var count int + require.NoError( + t, + f.pool.QueryRow(t.Context(), + `SELECT count(*) FROM agent_inputs WHERE agent_id=$1 AND input_kind='content'`, + a). + Scan(&count), + ) + require.Zero(t, count) + }) + } +} diff --git a/internal/storage/internal/agentexecution/commands_interactions.go b/internal/storage/internal/agentexecution/commands_interactions.go new file mode 100644 index 000000000..b98ef7f4d --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_interactions.go @@ -0,0 +1,459 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/interactionform" + "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/toolpermission" +) + +type OpenInteractionInput struct { + ToolRef + Permission *toolpermission.Request + Question *interactionform.Form + Destination json.RawMessage +} +type Interaction struct { + ID uuid.UUID + ToolID uuid.UUID + Kind string + State string + Request json.RawMessage + Resolution json.RawMessage + InputID uuid.UUID + Created bool +} +type ResolveInteractionInput struct { + ID uuid.UUID + ActorID uuid.UUID + Resolution interactionform.Resolution +} + +func interactionRecord(row executiondb.ReadExecutionInteractionRow) Interaction { + return Interaction{ID: row.ID, ToolID: row.ToolCallID, Kind: row.InteractionKind, State: row.State, + Request: row.Request, Resolution: row.Resolution, InputID: valueOrZero(row.ResolvedByInputID)} +} + +func interactionForm(kind string, request json.RawMessage) (interactionform.Form, error) { + if kind == "permission" { + r, err := toolpermission.ParseRequest(request) + return r.Form, err + } + return interactionform.Parse(request) +} + +func (h *Handle) OpenInteraction(ctx context.Context, input OpenInteractionInput) (Interaction, error) { + return executeCommand(ctx, h, func(m *executionMutation) (Interaction, error) { + if (input.Permission == nil) == (input.Question == nil) { + return Interaction{}, errors.New("exactly one interaction kind is required") + } + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return Interaction{}, err + } + kind := "question" + var request json.RawMessage + if input.Permission != nil { + kind = "permission" + if err := input.Permission.Validate(); err != nil { + return Interaction{}, err + } + request, _ = json.Marshal(input.Permission) + } else { + if err := input.Question.Validate(); err != nil { + return Interaction{}, err + } + request, _ = json.Marshal(input.Question) + } + q := executiondb.New() + existing, err := q.ReadExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ReadExecutionInteractionParams{ + AgentID: h.route.AgentID, ToolID: &input.ID, Kind: &kind}, + ) + if err == nil { + call, err := h.tool(ctx, input.ID) + if err != nil { + return Interaction{}, err + } + if kind == "permission" || call.Type != "built_in" || + call.State != "waiting" && call.State != "completed" || + !equalJSON(existing.Request, request) { + return Interaction{}, storeerr.ErrIdempotencyConflict + } + return interactionRecord(existing), nil + } + if !errors.Is(err, pgx.ErrNoRows) { + return Interaction{}, err + } + call, err := h.tool(ctx, input.ID) + if err != nil { + return Interaction{}, err + } + next := "waiting" + if kind == "permission" { + next = "awaiting_permission" + if call.State != "awaiting_authorization" { + return Interaction{}, storeerr.ErrStateTransitionConflict + } + } else if call.State != "ready" || call.Type != "built_in" { + return Interaction{}, storeerr.ErrStateTransitionConflict + } + var destination *json.RawMessage + if len(input.Destination) > 0 { + if !json.Valid(input.Destination) { + return Interaction{}, errors.New("invalid interaction destination") + } + destination = &input.Destination + } + id, err := q.InsertExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.InsertExecutionInteractionParams{ + AgentID: h.route.AgentID, ToolID: call.ID, Kind: kind, Request: request, Destination: destination}, + ) + if err != nil { + return Interaction{}, err + } + if err = h.transitionTool(ctx, m, call, next, uuid.Nil, input.RuntimeLockID); err != nil { + return Interaction{}, err + } + result := Interaction{ + ID: id, + ToolID: call.ID, + Kind: kind, + State: "open", + Request: request, + Resolution: json.RawMessage(`{}`), + Created: true, + } + h.interactionNotification(result, next) + if kind == "question" { + if err = h.questionEffect(ctx, result, *input.Question); err != nil { + return Interaction{}, err + } + } + return result, nil + }) +} + +func (h *Handle) interactionNotification(i Interaction, state string) { + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, i.ToolID, state, + ¬ifications.AgentInteractionUpdate{ID: i.ID, InteractionKind: i.Kind, State: i.State}) +} + +func (h *Handle) ResolveInteraction(ctx context.Context, input ResolveInteractionInput) (Interaction, error) { + return executeCommand(ctx, h, func(m *executionMutation) (Interaction, error) { + q := executiondb.New() + row, err := q.ReadExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ReadExecutionInteractionParams{AgentID: h.route.AgentID, ID: input.ID}, + ) + if err != nil { + return Interaction{}, err + } + form, err := interactionForm(row.InteractionKind, row.Request) + if err != nil { + return Interaction{}, err + } + normalized, err := interactionform.NormalizeResolution(form, input.Resolution) + if err != nil { + return Interaction{}, err + } + resolution, err := json.Marshal(normalized) + if err != nil { + return Interaction{}, err + } + if row.State != "open" { + if row.State != "resolved" || !equalJSON(row.Resolution, resolution) { + return Interaction{}, storeerr.ErrIdempotencyConflict + } + response, err := q.FindExecutionInput( + ctx, + h.unit.DB(), + executiondb.FindExecutionInputParams{ + AgentID: h.route.AgentID, + ID: valueOrZero(row.ResolvedByInputID), + }, + ) + if err != nil { + return Interaction{}, err + } + if valueOrZero(response.ActorID) != input.ActorID { + return Interaction{}, storeerr.ErrIdempotencyConflict + } + return interactionRecord(row), nil + } + call, err := h.tool(ctx, row.ToolCallID) + if err != nil { + return Interaction{}, err + } + if m.head.StopSequence > call.SourceSequence { + return Interaction{}, storeerr.ErrStateTransitionConflict + } + id, err := q.InsertExecutionResponse(ctx, h.unit.DB(), executiondb.InsertExecutionResponseParams{ + ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + ActorID: nullableID(input.ActorID), + InteractionID: &row.ID}) + if err != nil { + return Interaction{}, err + } + event, err := h.appendEvent(ctx, call.TurnID, "agent_input", id, uuid.Nil, uuid.Nil, false) + if err != nil { + return Interaction{}, err + } + if err = h.advanceTurn(ctx, event, true); err != nil { + return Interaction{}, err + } + n, err := q.ResolveExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ResolveExecutionInteractionParams{ + AgentID: h.route.AgentID, ID: row.ID, State: "resolved", Resolution: resolution, InputID: &id}, + ) + if err != nil { + return Interaction{}, err + } + if n != 1 { + return Interaction{}, storeerr.ErrStateTransitionConflict + } + n, err = q.ResolveExecutionInput( + ctx, + h.unit.DB(), + executiondb.ResolveExecutionInputParams{AgentID: h.route.AgentID, ID: id, EventID: &event.ID}, + ) + if err != nil { + return Interaction{}, err + } + if n != 1 { + return Interaction{}, storeerr.ErrStateTransitionConflict + } + result := interactionRecord(row) + result.State = "resolved" + result.Resolution = resolution + result.InputID = id + m.changed() + if err = h.applyInteraction(ctx, m, result, call); err != nil { + return Interaction{}, err + } + return result, nil + }) +} + +func (h *Handle) applyInteraction( + ctx context.Context, + m *executionMutation, + i Interaction, + call executiondb.ReadExecutionToolRow, +) error { + outcome := "succeeded" + var value any + if i.Kind == "permission" { + reason := "permission interaction canceled" + if i.State == "resolved" { + request, err := toolpermission.ParseRequest(i.Request) + if err != nil { + return err + } + answers, err := interactionform.ParseResolution(request.Form, i.Resolution) + if err != nil { + return err + } + decision, err := toolpermission.Resolve(request, answers) + if err != nil { + return err + } + if decision.Decision == toolpermission.DecisionAllow { + if call.State != "awaiting_permission" { + return storeerr.ErrStateTransitionConflict + } + if err = h.transitionTool(ctx, m, call, "ready", uuid.Nil, uuid.Nil); err != nil { + return err + } + h.interactionNotification(i, "ready") + return nil + } + reason = decision.Reason + } + outcome = "denied" + value = map[string]string{"reason": reason} + } else if i.State == "resolved" { + form, err := interactionform.Parse(i.Request) + if err != nil { + return err + } + answers, err := interactionform.ParseResolution(form, i.Resolution) + if err != nil { + return err + } + value = interactionform.RenderAnswers(form, answers) + } else { + value = map[string]string{"reason": "question interaction canceled"} + } + if i.State == "canceled" { + outcome = "canceled" + } + body, err := json.Marshal(value) + if err != nil { + return err + } + encoded, err := json.Marshal( + []map[string]any{{"type": "structured_data", "value": json.RawMessage(body)}}, + ) + if err != nil { + return err + } + if len(encoded) > 50*1024 { + outcome = "failed" + body = json.RawMessage( + `{"error":"Question response exceeded the size limit. Ask fewer or shorter questions, or request a shorter answer."}`, + ) + } + if _, + err = h.completeTool(ctx, + m, + call, + outcome, + []Content{{Kind: "structured_data", + Data: body}}, + uuid.Nil, + ¬ifications.AgentInteractionUpdate{ID: i.ID, + InteractionKind: i.Kind, + State: i.State}); err != nil { + return err + } + return nil +} + +func (h *Handle) SupersedeInteractions(ctx context.Context, inputID uuid.UUID) ([]uuid.UUID, error) { + return executeCommand(ctx, h, func(m *executionMutation) ([]uuid.UUID, error) { + row, err := executiondb.New(). + FindExecutionInput(ctx, + h.unit.DB(), + executiondb.FindExecutionInputParams{AgentID: h.route.AgentID, + ID: inputID}) + if err != nil { + return nil, err + } + if row.State != "received" || row.DeliveryMode != "steering" || row.InputKind != "content" { + return nil, storeerr.ErrStateTransitionConflict + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return nil, err + } + if !snapshot.Selection.TurnContinuable { + return nil, nil + } + return h.cancelInteractions(ctx, m, m.head.CurrentTurnID, inputID, "superseded_by_input") + }) +} + +func (h *Handle) cancelInteractions( + ctx context.Context, + m *executionMutation, + turn, input uuid.UUID, + reason string, +) ([]uuid.UUID, error) { + q := executiondb.New() + ids, err := q.ListExecutionInteractions( + ctx, + h.unit.DB(), + executiondb.ListExecutionInteractionsParams{AgentID: h.route.AgentID, TurnID: turn}, + ) + if err != nil { + return nil, err + } + resolution, err := json.Marshal(map[string]string{"reason": reason}) + if err != nil { + return nil, err + } + for _, id := range ids { + row, err := q.ReadExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ReadExecutionInteractionParams{AgentID: h.route.AgentID, ID: id}, + ) + if err != nil { + return nil, err + } + n, err := q.ResolveExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ResolveExecutionInteractionParams{ + AgentID: h.route.AgentID, ID: id, State: "canceled", Resolution: resolution, InputID: nullableID(input)}, + ) + if err != nil { + return nil, err + } + if n != 1 { + return nil, storeerr.ErrStateTransitionConflict + } + call, err := h.tool(ctx, row.ToolCallID) + if err != nil { + return nil, err + } + i := interactionRecord(row) + i.State = "canceled" + i.Resolution = resolution + i.InputID = input + if err = h.applyInteraction(ctx, m, i, call); err != nil { + return nil, err + } + } + return ids, nil +} + +func (h *Handle) RecordPresentation( + ctx context.Context, + id uuid.UUID, + destination, receipt json.RawMessage, +) (bool, error) { + if _, err := h.Route(); err != nil { + return false, err + } + if _, err := objectJSON(receipt); err != nil { + return false, err + } + q := executiondb.New() + row, err := q.ReadExecutionInteraction( + ctx, + h.unit.DB(), + executiondb.ReadExecutionInteractionParams{AgentID: h.route.AgentID, ID: id}, + ) + if err != nil { + return false, err + } + if row.Destination == nil || !equalJSON(*row.Destination, destination) { + return false, storeerr.ErrUnauthorized + } + if row.PresentationReceipt != nil { + if !equalJSON(*row.PresentationReceipt, receipt) { + return false, storeerr.ErrIdempotencyConflict + } + return false, nil + } + n, err := q.WriteExecutionPresentation( + ctx, + h.unit.DB(), + executiondb.WriteExecutionPresentationParams{ + AgentID: h.route.AgentID, + ID: id, + Destination: destination, + Receipt: receipt, + }, + ) + if err == nil && n != 1 { + return false, storeerr.ErrIdempotencyConflict + } + return n > 0, err +} diff --git a/internal/storage/internal/agentexecution/commands_lifecycle.go b/internal/storage/internal/agentexecution/commands_lifecycle.go new file mode 100644 index 000000000..5e17eb34f --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_lifecycle.go @@ -0,0 +1,305 @@ +package agentexecution + +import ( + "context" + "errors" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +func (h *Handle) cancelTurnProcesses(ctx context.Context, turn uuid.UUID) error { + q := dbsqlc.New(h.unit.DB()) + rows, err := q.CancelUnresolvedProcessesForAgentTurn( + ctx, + dbsqlc.CancelUnresolvedProcessesForAgentTurnParams{ + ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + TurnID: turn, + }, + ) + if err != nil { + return err + } + if _, + err = q.CancelQueuedProcessActionsForAgentTurn(ctx, + dbsqlc.CancelQueuedProcessActionsForAgentTurnParams{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + TurnID: turn}); err != nil { + return err + } + if _, + err = q.CancelAcceptedProcessActionsForAgentTurn(ctx, + dbsqlc.CancelAcceptedProcessActionsForAgentTurnParams{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + TurnID: turn}); err != nil { + return err + } + for _, row := range rows { + if err = h.interruptActions(ctx, + row.OrgID, + row.ID, + "agent_canceled_before_grant", + "agent_canceled_after_grant"); err != nil { + return err + } + if row.State == "unknown" { + h.unit.Notifications().AddDaemonProcessTermination(row.MachineID, row.ID) + } + } + return nil +} + +func (h *Handle) interruptActions( + ctx context.Context, + orgID, processID uuid.UUID, + queuedReason, acceptedReason string, +) error { + q := dbsqlc.New(h.unit.DB()) + queued, err := q.MarkQueuedProcessActionsFailedForProcess( + ctx, + dbsqlc.MarkQueuedProcessActionsFailedForProcessParams{ + OrgID: orgID, + ProcessID: processID, + StateReasonCode: optionalText(queuedReason), + }, + ) + if err != nil { + return err + } + accepted, err := q.ResolveAcceptedProcessActionsWithoutEvidence( + ctx, + dbsqlc.ResolveAcceptedProcessActionsWithoutEvidenceParams{ + OrgID: orgID, + ProcessID: processID, + StateReasonCode: optionalText(acceptedReason), + }, + ) + if err != nil { + return err + } + if len(accepted) > 0 { + if err = q.TouchProcessActivity(ctx, + dbsqlc.TouchProcessActivityParams{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + ProcessID: processID}); err != nil { + return err + } + } + for _, action := range append(queued, accepted...) { + if _, err = h.completeProcessAction(ctx, ActionResult{ID: action.ID}, true); err != nil { + return err + } + } + return nil +} + +func (h *Handle) failQueuedRuntime(ctx context.Context, id uuid.UUID, reason string) error { + q := dbsqlc.New(h.unit.DB()) + rows, err := q.FailQueuedProcessesForRuntimeEnd( + ctx, + dbsqlc.FailQueuedProcessesForRuntimeEndParams{ + ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + RuntimeLockID: id, + StateReasonCode: optionalText(reason), + }, + ) + if err != nil { + return err + } + for _, row := range rows { + if _, err = h.CompleteProcess(ctx, ProcessResult{ID: row.ID}); err != nil { + return err + } + } + actions, err := q.FailQueuedProcessActionsForRuntimeEnd( + ctx, + dbsqlc.FailQueuedProcessActionsForRuntimeEndParams{ + ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + RuntimeLockID: id, + StateReasonCode: optionalText(reason), + }, + ) + if err != nil { + return err + } + for _, row := range actions { + if _, err = h.completeProcessAction(ctx, ActionResult{ID: row.ID}, true); err != nil { + return err + } + } + return nil +} + +func (h *Handle) InterruptProcess(ctx context.Context, id uuid.UUID, reason string) (bool, error) { + return executeCommand(ctx, h, func(_ *executionMutation) (bool, error) { + if reason == "" { + return false, errors.New("process interruption reason is required") + } + row, err := executiondb.New(). + ReadExecutionProcessOwner(ctx, + h.unit.DB(), + executiondb.ReadExecutionProcessOwnerParams{AgentID: h.route.AgentID, + ID: id}) + if err != nil { + return false, err + } + q := dbsqlc.New(h.unit.DB()) + changed := false + switch row.State { + case "queued": + _, err = q.MarkQueuedProcessFailedByMachine( + ctx, + dbsqlc.MarkQueuedProcessFailedByMachineParams{ + OrgID: row.OrgID, + MachineID: row.MachineID, + ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + ID: id, + StateReasonCode: optionalText(reason), + }, + ) + if err != nil { + return false, err + } + changed = true + case "starting", "running": + _, err = q.MarkActiveProcessUnknownByMachine( + ctx, + dbsqlc.MarkActiveProcessUnknownByMachineParams{ + OrgID: row.OrgID, + MachineID: row.MachineID, + ID: id, + StateReasonCode: optionalText(reason), + }, + ) + if err != nil { + return false, err + } + changed = true + h.unit.Notifications().AddDaemonProcessTermination(row.MachineID, id) + } + if changed && row.ToolCallID != uuid.Nil { + if _, err = h.CompleteProcess(ctx, ProcessResult{ID: id}); err != nil { + return false, err + } + } + return changed, h.interruptActions(ctx, row.OrgID, id, reason, reason) + }) +} + +func (h *Handle) Archive(ctx context.Context, actorID uuid.UUID) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + identity, err := executiondb.New(). + ReadExecutionAgentIdentity(ctx, + h.unit.DB(), + executiondb.ReadExecutionAgentIdentityParams{ProjectID: h.route.ProjectID, + ID: h.route.AgentID}) + if err != nil { + return false, err + } + if identity.State == "archived" { + return false, nil + } + if err = h.unit.CaptureWebhookEligibility(ctx, h.route.AgentID); err != nil { + return false, err + } + n, err := executiondb.New(). + ArchiveExecutionAgent(ctx, h.unit.DB(), executiondb.ArchiveExecutionAgentParams{ID: h.route.AgentID}) + if err != nil { + return false, err + } + if n != 1 { + return false, storeerr.ErrStateTransitionConflict + } + q := dbsqlc.New(h.unit.DB()) + if err = q.DeleteAgentIntegrationSubscriptions(ctx, + dbsqlc.DeleteAgentIntegrationSubscriptionsParams{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID}); err != nil { + return false, err + } + if _, + err = q.DeleteCronTriggersForAgent(ctx, + dbsqlc.DeleteCronTriggersForAgentParams{ProjectID: h.route.ProjectID, + AgentID: &h.route.AgentID}); err != nil { + return false, err + } + rows, err := q.ListProcessesForExecutionRevoked( + ctx, + dbsqlc.ListProcessesForExecutionRevokedParams{ + ProjectID: h.route.ProjectID, + AgentID: &h.route.AgentID, + }, + ) + if err != nil { + return false, err + } + for _, row := range rows { + if _, err = h.InterruptProcess(ctx, row.ID, "agent_archived"); err != nil { + return false, err + } + } + if _, + err = h.Cancel(ctx, + CancelInput{ActorID: actorID, + Reason: "agent_archived", + Message: "The model call was stopped because the agent was archived.", + ForceRuntime: true}); err != nil { + return false, err + } + if _, + err = executiondb.New().CancelExecutionInputs(ctx, + h.unit.DB(), + executiondb.CancelExecutionInputsParams{AgentID: h.route.AgentID, + IncludeQueued: true}); err != nil { + return false, err + } + + m.changed() + return true, nil + }) +} + +func (u *Unit) ArchiveAgents(ctx context.Context, routes []AgentRoute, actorID uuid.UUID) error { + for _, route := range routes { + if _, err := u.Agent(route); err != nil { + return err + } + } + for _, route := range routes { + h := u.handles[route.AgentID] + if _, err := h.Archive(ctx, actorID); err != nil { + return err + } + } + return nil +} + +func (u *Unit) InterruptProcesses( + ctx context.Context, + agents map[AgentRoute][]uuid.UUID, + reason string, +) error { + for route := range agents { + if _, err := u.Agent(route); err != nil { + return err + } + } + for _, h := range u.orderedHandles() { + for _, id := range agents[h.route] { + _, err := h.InterruptProcess(ctx, id, reason) + if errors.Is(err, pgx.ErrNoRows) { + return storeerr.ErrNotFound + } + if err != nil { + return err + } + } + } + return nil +} diff --git a/internal/storage/internal/agentexecution/commands_models.go b/internal/storage/internal/agentexecution/commands_models.go new file mode 100644 index 000000000..1f9834c09 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_models.go @@ -0,0 +1,321 @@ +package agentexecution + +import ( + "context" + "errors" + "slices" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type PreparedModel struct { + Context ModelContext + ConfigID uuid.UUID + RevisionID uuid.UUID + RuntimeLockID uuid.UUID + Created bool + Claimed bool +} + +type PrepareModelInput struct { + RuntimeLockID uuid.UUID + Selected ModelDecision +} + +func (h *Handle) FenceRuntime(ctx context.Context, id uuid.UUID) error { + if id == uuid.Nil { + return storeerr.ErrRuntimeLockInactive + } + if err := h.unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID}}, + RuntimeAuthority{AgentID: h.route.AgentID, RuntimeLockID: id}); err != nil { + return err + } + _, err := executiondb.New().FenceExecutionRuntime(ctx, h.unit.DB(), executiondb.FenceExecutionRuntimeParams{ + ProjectID: h.route.ProjectID, AgentID: h.route.AgentID, RuntimeID: id}) + if errors.Is(err, pgx.ErrNoRows) { + return storeerr.ErrRuntimeLockInactive + } + + return err +} + +func attemptRecord(row executiondb.ReadExecutionAttemptRow) PreparedModel { + return PreparedModel{Context: ModelContext{ID: row.ID, AgentID: row.AgentID, TurnID: row.TurnID, + Operation: Operation(row.OperationKind), + Attempt: row.AttemptNumber, + InputEventSequence: row.InputEventSequence, + SourceEventSequenceEnd: valueOrZero(row.SourceEventSequenceEnd), State: ContextState(row.State), + Recovery: RecoveryKind(valueOrZero(row.RecoveryKind)), RetryAt: row.RetryAt, + Opening: Opening{InputIDs: row.OpeningInputIds, EventSequence: row.OpeningEventSequence}}, + ConfigID: row.AgentConfigID, RevisionID: row.ConfiguredModelRevisionID, RuntimeLockID: row.RuntimeLockID} +} + +func sameDecision(a, b ModelDecision) bool { + return a.Kind == b.Kind && + a.Origin == b.Origin && + a.TurnID == b.TurnID && + a.SourceContextID == b.SourceContextID && + a.SourceOutputID == b.SourceOutputID && + a.SourceInputID == b.SourceInputID && + a.SourceCheckpointID == b.SourceCheckpointID && + a.Opening.EventSequence == b.Opening.EventSequence && slices.Equal(a.Opening.InputIDs, b.Opening.InputIDs) +} + +func (h *Handle) PrepareModel(ctx context.Context, input PrepareModelInput) (PreparedModel, error) { + return executeCommand(ctx, h, func(m *executionMutation) (PreparedModel, error) { + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return PreparedModel{}, err + } + q := executiondb.New() + boundary, err := q.ExecutionPreparationBoundary( + ctx, + h.unit.DB(), + executiondb.ExecutionPreparationBoundaryParams{AgentID: h.route.AgentID}, + ) + if err != nil { + return PreparedModel{}, err + } + desired := PreparedModel{ConfigID: boundary.CurrentConfigID, RuntimeLockID: input.RuntimeLockID, + Context: ModelContext{ + ID: uuid.New(), + AgentID: h.route.AgentID, + TurnID: input.Selected.TurnID, + Operation: OperationNormal, + Attempt: 1, + InputEventSequence: boundary.Watermark, + State: ContextStarted, + Opening: input.Selected.Opening, + }} + if input.Selected.Kind == ModelResume { + source, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: input.Selected.SourceContextID}, + ) + if err != nil { + return PreparedModel{}, err + } + desired = attemptRecord(source) + desired.Context.ID = uuid.New() + desired.Context.Attempt++ + desired.Context.State = ContextStarted + desired.Context.Recovery = RecoveryNone + desired.Context.RetryAt = nil + desired.RuntimeLockID = input.RuntimeLockID + } + id, err := q.LatestExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.LatestExecutionAttemptParams{AgentID: h.route.AgentID, + Operation: string(desired.Context.Operation), + Watermark: desired.Context.InputEventSequence, + SourceEnd: optionalSequence(desired.Context.SourceEventSequenceEnd)}, + ) + if err == nil { + existing, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: id}, + ) + if err != nil { + return PreparedModel{}, err + } + record := attemptRecord(existing) + if record.Context.Attempt >= desired.Context.Attempt { + if record.Context.TurnID != input.Selected.TurnID || record.ConfigID != desired.ConfigID || + !slices.Equal(record.Context.Opening.InputIDs, input.Selected.Opening.InputIDs) || + record.Context.Opening.EventSequence != input.Selected.Opening.EventSequence { + return PreparedModel{}, storeerr.ErrIdempotencyConflict + } + return record, nil + } + } else if !errors.Is(err, pgx.ErrNoRows) { + return PreparedModel{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return PreparedModel{}, err + } + candidate := snapshot.Selection.Model + if candidate == nil || !sameDecision(*candidate, input.Selected) || boundary.Started || + (input.Selected.Kind == ModelResume && (candidate.ReadyAt.After(boundary.DatabaseNow) || + desired.Context.Attempt > 9)) { + return PreparedModel{}, storeerr.ErrAgentNotAdvanceable + } + return h.createAttempt(ctx, m, desired, snapshot) + }) +} + +func (h *Handle) createAttempt( + ctx context.Context, + m *executionMutation, + desired PreparedModel, + snapshot ExecutionSnapshot, +) (PreparedModel, error) { + q := executiondb.New() + revision, err := q.ExecutionModelRevision( + ctx, + h.unit.DB(), + executiondb.ExecutionModelRevisionParams{ProjectID: h.route.ProjectID, ConfigID: desired.ConfigID}, + ) + if err != nil { + return PreparedModel{}, err + } + c := desired.Context + _, err = q.CreateExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.CreateExecutionAttemptParams{ID: c.ID, AgentID: h.route.AgentID, + ProjectID: h.route.ProjectID, + TurnID: c.TurnID, + Operation: string(c.Operation), + Attempt: c.Attempt, + ConfigID: desired.ConfigID, + RevisionID: revision.CurrentRevisionID, + Watermark: c.InputEventSequence, + SourceEnd: optionalSequence(c.SourceEventSequenceEnd), + RuntimeID: desired.RuntimeLockID, + OpeningIds: c.Opening.InputIDs, + OpeningSequence: c.Opening.EventSequence}, + ) + if errors.Is(err, pgx.ErrNoRows) { + return PreparedModel{}, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return PreparedModel{}, err + } + m.head.MaxContextInputSequence = max(m.head.MaxContextInputSequence, c.InputEventSequence) + if c.Operation == OperationNormal { + m.head.MaxNormalInputSequence = max(m.head.MaxNormalInputSequence, c.InputEventSequence) + m.head.NormalContextID = c.ID + m.head.CompactionContextID = uuid.Nil + if snapshot.View.Config != nil && snapshot.View.Config.Event.Sequence <= c.InputEventSequence { + m.head.PendingConfigInputID = uuid.Nil + } + if snapshot.View.Checkpoint != nil && snapshot.View.Checkpoint.Event.Sequence <= c.InputEventSequence { + m.head.PendingCheckpointID = uuid.Nil + } + if snapshot.View.OutputLimit != nil && snapshot.View.OutputLimit.Event.Sequence <= c.InputEventSequence { + m.head.PendingOutputLimitID = uuid.Nil + } + } else { + m.head.CompactionContextID = c.ID + } + m.changed() + desired.RevisionID = revision.CurrentRevisionID + desired.Created = true + desired.Claimed = revision.Allowed + if !revision.Allowed { + _, err := h.terminalFailure(ctx, m, ModelFailure{ContextID: c.ID, RuntimeLockID: desired.RuntimeLockID, + ErrorKind: "runtime", + ErrorCode: storeerr.ManagedWorkAdmissionDeniedCode, + ErrorMessage: storeerr.InsufficientOmnaraCreditsMessage}) + if err != nil { + return PreparedModel{}, err + } + desired.Context.State = ContextFailed + } + return desired, nil +} + +type PrepareNormalInput struct { + RuntimeLockID uuid.UUID + OpeningInputIDs []uuid.UUID + SourceContextID uuid.UUID + SourceOutputID uuid.UUID +} + +func (h *Handle) PrepareNormal(ctx context.Context, input PrepareNormalInput) (PreparedModel, error) { + return executeCommand(ctx, h, func(_ *executionMutation) (PreparedModel, error) { + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return PreparedModel{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return PreparedModel{}, err + } + selected := snapshot.Selection.Model + if selected == nil { + governing := snapshot.View.NormalContext + if governing == nil { + return PreparedModel{}, storeerr.ErrAgentNotAdvanceable + } + selected = &ModelDecision{ + Kind: ModelStart, + Origin: OriginInitial, + TurnID: governing.TurnID, + Opening: governing.Opening, + } + } + if !slices.Equal(selected.Opening.InputIDs, input.OpeningInputIDs) || + selected.Kind == ModelResume || + (selected.Kind == ModelContinue && + (selected.SourceContextID != input.SourceContextID || + selected.SourceOutputID != input.SourceOutputID)) || + (selected.Kind != ModelContinue && + (input.SourceContextID != uuid.Nil || + input.SourceOutputID != uuid.Nil)) { + return PreparedModel{}, storeerr.ErrAgentNotAdvanceable + } + prepared, err := h.PrepareModel( + ctx, + PrepareModelInput{RuntimeLockID: input.RuntimeLockID, Selected: *selected}, + ) + if err != nil { + return PreparedModel{}, err + } + return prepared, nil + }) +} + +func (h *Handle) ResumeModel(ctx context.Context, runtimeID, predecessorID uuid.UUID) (PreparedModel, error) { + return executeCommand(ctx, h, func(_ *executionMutation) (PreparedModel, error) { + if err := h.FenceRuntime(ctx, runtimeID); err != nil { + return PreparedModel{}, err + } + q := executiondb.New() + predecessor, err := q.ReadExecutionAttempt(ctx, h.unit.DB(), executiondb.ReadExecutionAttemptParams{ + AgentID: h.route.AgentID, ID: predecessorID, + }) + if err != nil { + return PreparedModel{}, err + } + id, err := q.LatestExecutionAttempt(ctx, h.unit.DB(), executiondb.LatestExecutionAttemptParams{ + AgentID: h.route.AgentID, Operation: predecessor.OperationKind, + Watermark: predecessor.InputEventSequence, SourceEnd: predecessor.SourceEventSequenceEnd, + }) + if err != nil { + return PreparedModel{}, err + } + latest, err := q.ReadExecutionAttempt(ctx, h.unit.DB(), executiondb.ReadExecutionAttemptParams{ + AgentID: h.route.AgentID, ID: id, + }) + if err != nil { + return PreparedModel{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return PreparedModel{}, err + } + selected := snapshot.Selection.Model + if selected == nil || selected.Kind != ModelResume || selected.SourceContextID != predecessor.ID { + return attemptRecord(latest), nil + } + boundary, err := q.ExecutionPreparationBoundary(ctx, h.unit.DB(), executiondb.ExecutionPreparationBoundaryParams{ + AgentID: h.route.AgentID, + }) + if err != nil { + return PreparedModel{}, err + } + if boundary.Started || selected.ReadyAt.After(boundary.DatabaseNow) || predecessor.AttemptNumber >= 9 { + return attemptRecord(latest), nil + } + return h.PrepareModel(ctx, PrepareModelInput{RuntimeLockID: runtimeID, Selected: *selected}) + }) +} diff --git a/internal/storage/internal/agentexecution/commands_outputs.go b/internal/storage/internal/agentexecution/commands_outputs.go new file mode 100644 index 000000000..43e0e5105 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_outputs.go @@ -0,0 +1,369 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "strings" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ToolProposal struct { + ID uuid.UUID + ProviderCallID string + Type string +} +type AcceptOutputInput struct { + RuntimeLockID uuid.UUID + ContextID uuid.UUID + RequestID string + Response modelenvelope.ResponseEnvelope + Tools []ToolProposal +} + +type AcceptedOutput struct { + ID uuid.UUID + Event ExecutionEvent + Tools []ToolProposal + Created bool +} + +type outputProposal struct { + ToolProposal + Name string + Input json.RawMessage + Error string +} + +func outputParts( + response modelenvelope.ResponseEnvelope, + bindings []ToolProposal, +) ([]Content, []outputProposal, error) { + byID := make(map[string]ToolProposal, len(bindings)) + for _, binding := range bindings { + if binding.ProviderCallID == "" || + binding.Type != "built_in" && binding.Type != "custom" && binding.Type != "mcp" { + return nil, nil, errors.New("invalid tool proposal") + } + if _, exists := byID[binding.ProviderCallID]; exists { + return nil, nil, errors.New("duplicate tool proposal") + } + if binding.ID == uuid.Nil { + binding.ID = uuid.New() + } + byID[binding.ProviderCallID] = binding + } + var content []Content + var proposals []outputProposal + for ordinal, part := range response.Normalized.Content { + switch part.Type { + case modelenvelope.ResponsePartTypeText, + modelenvelope.ResponsePartTypeReasoning, + modelenvelope.ResponsePartTypeError: + if part.Text != "" { + content = append( + content, + Content{Ordinal: int32(ordinal), Kind: string(part.Type), Text: part.Text}, + ) + } + case modelenvelope.ResponsePartTypeToolCall: + binding, found := byID[part.ProviderCallID] + if !found { + return nil, nil, errors.New("provider tool call has no tool call binding") + } + if err := modelenvelope.ValidateToolInput(part.ToolInput); err != nil { + return nil, nil, err + } + delete(byID, part.ProviderCallID) + proposals = append( + proposals, + outputProposal{ + ToolProposal: binding, + Name: part.ToolName, + Input: part.ToolInput, + Error: part.ToolCallError, + }, + ) + content = append(content, Content{Ordinal: int32(ordinal), Kind: "tool_call", ToolID: binding.ID}) + default: + return nil, nil, errors.New("unsupported output content") + } + } + if len(byID) > 0 { + return nil, nil, errors.New("unused tool bindings") + } + return content, proposals, nil +} + +func (h *Handle) AcceptOutput(ctx context.Context, input AcceptOutputInput) (AcceptedOutput, error) { + return executeCommand(ctx, h, func(m *executionMutation) (AcceptedOutput, error) { + if err := input.Response.Validate(); err != nil { + return AcceptedOutput{}, err + } + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return AcceptedOutput{}, err + } + q := executiondb.New() + source, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: input.ContextID}, + ) + if err != nil { + return AcceptedOutput{}, err + } + if source.OperationKind != "normal" || + source.ProviderModelSlug != input.Response.RequestedProviderModelSlug { + return AcceptedOutput{}, storeerr.ErrStateTransitionConflict + } + if source.RuntimeLockID != input.RuntimeLockID { + return AcceptedOutput{}, storeerr.ErrRuntimeLockInactive + } + if source.State != "started" && source.State != "succeeded" { + return AcceptedOutput{}, storeerr.ErrStateTransitionConflict + } + evidence := ModelEvidence{APIFormat: input.Response.APIFormat, APIVariant: input.Response.APIVariant, + RequestID: input.RequestID, + ResponseID: input.Response.Normalized.ID, + Usage: input.Response.Normalized.Usage, + Cost: input.Response.ProviderReportedCostUSD, Metadata: input.Response.ProviderMetadata} + if source.State == "succeeded" { + return h.replayOutput(ctx, m, source, input, evidence) + } + parts, proposals, err := outputParts(input.Response, input.Tools) + if err != nil { + return AcceptedOutput{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return AcceptedOutput{}, err + } + var replay *json.RawMessage + if len(input.Response.ProviderReplay) > 0 { + replay = &input.Response.ProviderReplay + } + output, err := q.CreateExecutionOutput(ctx, h.unit.DB(), executiondb.CreateExecutionOutputParams{ + AgentID: h.route.AgentID, ContextID: input.ContextID, ServedModel: input.Response.ServedProviderModelSlug, + StopReason: string(input.Response.Normalized.StopReason), Replay: replay}) + if err != nil { + return AcceptedOutput{}, err + } + event, err := h.appendEvent( + ctx, + source.TurnID, + "model_output", + uuid.Nil, + output.ID, + uuid.Nil, + false, + ) + if err != nil { + return AcceptedOutput{}, err + } + result := AcceptedOutput{ID: output.ID, Event: event, Created: true} + for _, call := range proposals { + if err = q.CreateExecutionToolProposal(ctx, h.unit.DB(), executiondb.CreateExecutionToolProposalParams{ + ID: call.ID, AgentID: h.route.AgentID, OutputID: output.ID, ProviderCallID: call.ProviderCallID, + Name: call.Name, Input: call.Input, Type: call.Type}); err != nil { + return AcceptedOutput{}, err + } + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, call.ID, "awaiting_authorization", nil) + result.Tools = append(result.Tools, call.ToolProposal) + } + if err = h.writeContent(ctx, "model_output", output.ID, parts); err != nil { + return AcceptedOutput{}, err + } + if err = h.advanceTurn(ctx, event, len(proposals) == 0); err != nil { + return AcceptedOutput{}, err + } + for _, call := range proposals { + if call.Error != "" { + if err = h.rejectProposal(ctx, m, call); err != nil { + return AcceptedOutput{}, err + } + } + } + if _, err = h.finishAttempt( + ctx, m, source, ContextSucceeded, ModelFailure{RuntimeLockID: input.RuntimeLockID}, evidence, + ); err != nil { + return AcceptedOutput{}, err + } + if err = applyAcceptedOutput(m, + snapshot, + attemptRecord(source).Context, + output.ID, + event, + len(proposals) > 0, + input.Response.Normalized.StopReason == modelenvelope.StopReasonMaxTokens); err != nil { + return AcceptedOutput{}, err + } + if !input.Response.HasToolCalls() { + kind := map[modelenvelope.StopReason]string{modelenvelope.StopReasonEndTurn: "result", + modelenvelope.StopReasonRefusal: "refused", + modelenvelope.StopReasonContentFilter: "content_filtered"}[input.Response.Normalized.StopReason] + if kind != "" { + if err = h.parentEffect(ctx, + kind, + strings.TrimSpace(input.Response.Text()), + "model_output:"+output.ID.String(), + uuid.Nil); err != nil { + return AcceptedOutput{}, err + } + } + } + return result, nil + }) +} + +func (h *Handle) replayOutput( + ctx context.Context, + m *executionMutation, + source executiondb.ReadExecutionAttemptRow, + input AcceptOutputInput, + evidence ModelEvidence, +) (AcceptedOutput, error) { + if _, err := h.finishAttempt( + ctx, m, source, ContextSucceeded, ModelFailure{RuntimeLockID: input.RuntimeLockID}, evidence, + ); err != nil { + return AcceptedOutput{}, err + } + q := executiondb.New() + output, err := q.FindExecutionOutput( + ctx, + h.unit.DB(), + executiondb.FindExecutionOutputParams{AgentID: h.route.AgentID, ContextID: source.ID}, + ) + if err != nil { + return AcceptedOutput{}, err + } + replayPresent := len(input.Response.ProviderReplay) > 0 || output.ProviderReplay != nil + replayMatches := equalJSON(valueOrZero(output.ProviderReplay), input.Response.ProviderReplay) + if output.StopReason != string(input.Response.Normalized.StopReason) || + output.ServedProviderModelSlug != input.Response.ServedProviderModelSlug || + (replayPresent && !replayMatches) { + return AcceptedOutput{}, storeerr.ErrIdempotencyConflict + } + stored, err := q.ReadExecutionProposals( + ctx, + h.unit.DB(), + executiondb.ReadExecutionProposalsParams{AgentID: h.route.AgentID, OutputID: output.ID}, + ) + if err != nil { + return AcceptedOutput{}, err + } + bindings := make([]ToolProposal, len(input.Tools)) + copy(bindings, input.Tools) + for i, binding := range bindings { + for _, call := range stored { + if binding.ProviderCallID == call.ProviderCallID && binding.ID == uuid.Nil { + bindings[i].ID = call.ID + } + } + } + parts, proposals, err := outputParts(input.Response, bindings) + if err != nil { + return AcceptedOutput{}, err + } + if len(stored) != len(proposals) { + return AcceptedOutput{}, storeerr.ErrIdempotencyConflict + } + for _, call := range stored { + matched := false + for _, proposal := range proposals { + if proposal.ID == call.ID && proposal.ProviderCallID == call.ProviderCallID && + proposal.Name == call.Name && + proposal.Type == call.Type && + equalJSON(proposal.Input, call.Input) { + if proposal.Error != "" { + row, err := h.tool(ctx, call.ID) + if err != nil { + return AcceptedOutput{}, err + } + body, err := json.Marshal( + map[string]string{"error": proposal.Error, "error_code": "malformed"}, + ) + if err != nil { + return AcceptedOutput{}, err + } + if row.ResultID == nil || valueOrZero(row.Outcome) != "failed" { + return AcceptedOutput{}, storeerr.ErrIdempotencyConflict + } + if err = h.matchContent(ctx, + executiondb.ReadExecutionContentParams{AgentID: h.route.AgentID, + ResultID: row.ResultID}, + []Content{{Kind: "structured_data", + Data: body}}); err != nil { + return AcceptedOutput{}, err + } + } + matched = true + break + } + } + if !matched { + return AcceptedOutput{}, storeerr.ErrIdempotencyConflict + } + } + if err = h.contentMatches(ctx, uuid.Nil, output.ID, parts); err != nil { + return AcceptedOutput{}, err + } + return AcceptedOutput{ID: output.ID, Event: ExecutionEvent{ID: output.EventID, TurnID: output.TurnID, + EventBoundary: EventBoundary{ + Sequence: output.Sequence, + Time: output.CreatedAt, + }}, Tools: bindings}, nil +} + +func applyAcceptedOutput( + m *executionMutation, + snapshot ExecutionSnapshot, + source ModelContext, + id uuid.UUID, + event ExecutionEvent, + tools, limit bool, +) error { + m.head.AnsweredThroughSequence = max(m.head.AnsweredThroughSequence, event.Sequence) + m.changed() + if source.TurnID != m.head.CurrentTurnID || event.Sequence < m.head.StopSequence { + return nil + } + if snapshot.View.ToolBatch != nil { + completion := snapshot.View.ToolBatch.Completion + if completion == nil || source.InputEventSequence < completion.LastResultSequence { + return errors.New("accepted model output does not cover pending tool results") + } + m.head.PendingToolOutputID = uuid.Nil + } + m.head.PendingOutputLimitID = uuid.Nil + if tools { + m.head.PendingToolOutputID = id + } else if limit { + m.head.PendingOutputLimitID = id + } + return nil +} + +func (h *Handle) rejectProposal(ctx context.Context, m *executionMutation, call outputProposal) error { + row, err := h.tool(ctx, call.ID) + if err != nil { + return err + } + data, err := json.Marshal(map[string]string{"error": call.Error, "error_code": "malformed"}) + if err != nil { + return err + } + _, err = h.completeTool( + ctx, + m, + row, + "failed", + []Content{{Kind: "structured_data", Data: data}}, + uuid.Nil, + nil, + ) + return err +} diff --git a/internal/storage/internal/agentexecution/commands_permission_integration_test.go b/internal/storage/internal/agentexecution/commands_permission_integration_test.go new file mode 100644 index 000000000..af6cc5fde --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_permission_integration_test.go @@ -0,0 +1,144 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/interactionform" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/toolpermission" + "github.com/stretchr/testify/require" +) + +func TestExecutionPermissionAndPresentation(t *testing.T) { + for _, mode := range []string{"allow", "deny", "supersede", "cancel"} { + t.Run(mode, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + ref := toolCommand(t, h, lease, "built_in") + form, err := toolpermission.NewAllowDenyForm("Permission", nil) + require.NoError(t, err) + permission := toolpermission.Request{ + Permission: toolpermission.Selection{Mode: toolpermission.ModeAlwaysAsk}, + Authorization: toolpermission.Authorization{ToolName: "tool", Input: json.RawMessage(`{}`)}, + Form: form, + } + dest := json.RawMessage(`{"target":"test"}`) + opened, err := h.OpenInteraction( + t.Context(), + agentexecution.OpenInteractionInput{ToolRef: ref, Permission: &permission, Destination: dest}, + ) + require.NoError(t, err) + assertCommandState(t, u, h) + if mode == "cancel" { + _, err = h.Cancel( + t.Context(), + agentexecution.CancelInput{Reason: "agent_canceled", Message: "canceled"}, + ) + require.NoError(t, err) + } else if mode == "supersede" { + received := receiveCommand(t, h, "steering", "supersede") + ids, err := h.SupersedeInteractions(t.Context(), received.ID) + require.NoError(t, err) + require.Equal(t, []uuid.UUID{opened.ID}, ids) + } else { + option := toolpermission.AllowOptionIndex + if mode == "deny" { + option = toolpermission.DenyOptionIndex + } + input := agentexecution.ResolveInteractionInput{ID: opened.ID, + Resolution: interactionform.Resolution{Answers: []interactionform.Answer{{OptionIndices: []int{option}}}}} + resolved, err := h.ResolveInteraction(t.Context(), input) + require.NoError(t, err) + replay, err := h.ResolveInteraction(t.Context(), input) + require.NoError(t, err) + require.Equal(t, resolved.InputID, replay.InputID) + } + receipt := json.RawMessage(`{"message":"sent"}`) + changed, err := h.RecordPresentation(t.Context(), opened.ID, dest, receipt) + require.NoError(t, err) + require.True(t, changed) + changed, err = h.RecordPresentation(t.Context(), opened.ID, dest, receipt) + require.NoError(t, err) + require.False(t, changed) + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.RecordPresentation(t.Context(), opened.ID, json.RawMessage(`{}`), receipt) + return err + }), storeerr.ErrUnauthorized) + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.RecordPresentation(t.Context(), opened.ID, dest, json.RawMessage(`{}`)) + return err + }), storeerr.ErrIdempotencyConflict) + snapshot := assertCommandState(t, u, h) + if mode == "allow" { + require.Equal(t, agentexecution.WorkTool, snapshot.Selection.Work) + } + require.NoError(t, u.Commit(t.Context(), "permission")) + }) + } +} + +func TestExecutionRepairAndMetadata(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "metadata") + changed, err := h.SetInteractionTarget(t.Context(), uuid.Nil, "", true) + require.NoError(t, err) + require.True(t, changed) + _, err = h.Repair(t.Context()) + require.NoError(t, err) + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "repair metadata")) + u, h = f.handle(t, a) + _, err = u.DB().Exec(t.Context(), `DELETE FROM agent_execution_state WHERE agent_id=$1`, a) + require.NoError(t, err) + _, err = h.Repair(t.Context()) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "repair missing head")) + u, h = f.handle(t, a) + assertCommandState(t, u, h) + require.NoError(t, u.ClearIntegrationTargets(t.Context(), f.project, uuid.New())) + require.NoError(t, u.Commit(t.Context(), "empty integration cleanup")) +} + +func TestExecutionIntegrationTargetCleanup(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + integration, target := uuid.New(), uuid.New() + _, err := u.DB().Exec(t.Context(), `INSERT INTO integrations(id,org_id,project_id,name,integration_kind,settings, + state,provider_tenant_id,provider_account_ref,credential_secret_id,installed_by_user_id,created_at,updated_at) + VALUES($1,$2,$3,'test','slack_thread','{}','active','tenant','account', + (SELECT id FROM secrets WHERE org_id=$2 LIMIT 1),(SELECT id FROM users LIMIT 1), + statement_timestamp(),statement_timestamp())`, integration, f.org, f.project) + require.NoError(t, err) + _, err = u.DB().Exec(t.Context(), `INSERT INTO integration_targets(id,project_id,agent_id,integration_id, + scope_kind,scope_ref,created_at,updated_at) VALUES($1,$2,$3,$4,'thread','C123:111.222', + statement_timestamp(),statement_timestamp())`, target, f.project, a, integration) + require.NoError(t, err) + changed, err := h.SetInteractionTarget(t.Context(), target, "handler", false) + require.NoError(t, err) + require.True(t, changed) + _, + err = u.DB().Exec(t.Context(), + `UPDATE integration_targets SET deleted_at=statement_timestamp() WHERE id=$1`, + target) + require.NoError(t, err) + require.NoError(t, u.ClearIntegrationTargets(t.Context(), f.project, integration)) + var selected *uuid.UUID + var automatic bool + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT interaction_target_id,interaction_auto_select FROM agents WHERE id=$1`, + a).Scan(&selected, + &automatic)) + require.Nil(t, selected) + require.False(t, automatic) + statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) + require.NoError(t, err) + require.Zero(t, statements) + require.Zero(t, roundTrips) +} diff --git a/internal/storage/internal/agentexecution/commands_processes.go b/internal/storage/internal/agentexecution/commands_processes.go new file mode 100644 index 000000000..070a68b75 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_processes.go @@ -0,0 +1,305 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/processresult" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ProcessResult struct { + ID uuid.UUID + Started bool + Observed json.RawMessage + Content []Content + Outcome string +} + +type ActionResult struct { + ID uuid.UUID + Observed json.RawMessage +} + +type ExternalCompletion struct { + Tool CompletedTool + Matches bool + Published bool +} + +func processFacts(row executiondb.ReadExecutionProcessOwnerRow) processresult.Process { + var exit *int + if row.ExitCode != nil { + n := int(*row.ExitCode) + exit = &n + } + return processresult.Process{ID: row.ID, State: row.State, DefaultOutputCursor: row.DefaultOutputCursor, + SourceStartedAt: row.SourceStartedAt, + SourceEndedAt: row.SourceEndedAt, + ExitCode: exit, + ExitSignal: row.ExitSignal, + StateReasonCode: valueOrZero( + row.StateReasonCode, + ), StateReasonMessage: row.StateReasonMessage, ExecutionSpec: row.ExecutionSpec} +} + +func (h *Handle) CompleteProcess(ctx context.Context, input ProcessResult) (ExternalCompletion, error) { + return executeCommand(ctx, h, func(m *executionMutation) (ExternalCompletion, error) { + row, err := executiondb.New(). + ReadExecutionProcessOwner(ctx, + h.unit.DB(), + executiondb.ReadExecutionProcessOwnerParams{AgentID: h.route.AgentID, + ID: input.ID}) + if err != nil { + return ExternalCompletion{}, err + } + call, err := h.tool(ctx, row.ToolCallID) + if err != nil { + return ExternalCompletion{}, err + } + outcome := "succeeded" + var result json.RawMessage + if input.Started { + if row.ExecutionSpec.FileTransfer != nil { + return ExternalCompletion{}, storeerr.ErrInvalidToolCallDisposition + } + if row.State != "running" && call.State != "completed" { + return ExternalCompletion{}, storeerr.ErrStateTransitionConflict + } + facts := processFacts(row) + facts.State = "running" + result, err = processresult.Started(facts, input.Observed) + } else { + if !processresult.TerminalState(row.State) { + return ExternalCompletion{}, storeerr.ErrStateTransitionConflict + } + outcome, result, err = processresult.Terminal(processFacts(row)) + if err == nil && len(input.Observed) > 0 { + var body map[string]any + err = json.Unmarshal(input.Observed, &body) + if err == nil && body == nil { + err = errors.New("terminal process observation must be an object") + } + if err == nil { + var base map[string]any + err = json.Unmarshal(result, &base) + if err == nil { + body["process_id"] = base["process_id"] + result, err = json.Marshal(body) + } + } + } + } + if err != nil { + return ExternalCompletion{}, err + } + parts := []Content{{Kind: "structured_data", Data: result}} + if input.Content != nil { + if row.ExecutionSpec.FileTransfer == nil || input.Started || + input.Outcome != "succeeded" && input.Outcome != "failed" { + return ExternalCompletion{}, storeerr.ErrInvalidToolCallDisposition + } + outcome = input.Outcome + parts = input.Content + } + published, err := h.completeExternalTool(ctx, m, call, outcome, parts) + if err != nil { + return ExternalCompletion{}, err + } + if !published.Matches { + return published, nil + } + next, found, err := processresult.ObservationNextCursor(result) + if err != nil { + return ExternalCompletion{}, err + } + if published.Matches && found { + _, err = dbsqlc.New(h.unit.DB()). + AdvanceProcessDefaultOutputCursor(ctx, + dbsqlc.AdvanceProcessDefaultOutputCursorParams{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + ID: row.ID, + NextCursor: next}) + } + return published, err + }) +} + +func (h *Handle) CompleteProcessAction(ctx context.Context, input ActionResult) (ExternalCompletion, error) { + return h.completeProcessAction(ctx, input, false) +} + +func (h *Handle) SettleInterruptedAction(ctx context.Context, id uuid.UUID) (ExternalCompletion, error) { + return h.completeProcessAction(ctx, ActionResult{ID: id}, true) +} + +func (h *Handle) completeProcessAction( + ctx context.Context, + input ActionResult, + interrupted bool, +) (ExternalCompletion, error) { + return executeCommand(ctx, h, func(m *executionMutation) (ExternalCompletion, error) { + q := executiondb.New() + action, err := q.ReadExecutionActionOwner( + ctx, + h.unit.DB(), + executiondb.ReadExecutionActionOwnerParams{AgentID: h.route.AgentID, ID: input.ID}, + ) + if err != nil { + return ExternalCompletion{}, err + } + if action.State != "applied" && action.State != "failed" && action.State != "unknown" { + return ExternalCompletion{}, storeerr.ErrStateTransitionConflict + } + process, err := q.ReadExecutionProcessOwner( + ctx, + h.unit.DB(), + executiondb.ReadExecutionProcessOwnerParams{AgentID: h.route.AgentID, ID: action.ProcessID}, + ) + if err != nil { + return ExternalCompletion{}, err + } + call, err := h.tool(ctx, action.ToolCallID) + if err != nil { + return ExternalCompletion{}, err + } + if !interrupted && action.EarlierPending && call.State != "completed" && action.State == "applied" { + return ExternalCompletion{}, storeerr.ErrProcessActionReportBlocked + } + outcome := "succeeded" + if action.State != "applied" { + outcome = "failed" + } + result := input.Observed + var advance *int64 + if action.ActionKind == "read" { + facts := processFacts(process) + request := processresult.Action{ID: action.ID, Payload: action.Payload} + if action.State == "applied" { + if call.State == "completed" { + return h.replayProcessRead(ctx, call, result) + } + result, advance, err = processresult.Read(facts, request, result) + } else { + message := action.StateReasonMessage + if interrupted { + message = valueOrZero(action.StateReasonCode) + } + result, err = processresult.ReadFailure(facts, request, valueOrZero(action.StateReasonCode), message) + } + } else if len(result) == 0 || string(result) == "null" { + message := action.StateReasonMessage + if message == "" && action.State != "applied" { + message = valueOrZero(action.StateReasonCode) + } + result, + err = processresult.ActionResult(process.ID, + action.ID, + action.State, + valueOrZero(action.StateReasonCode), + message) + } + if err != nil { + return ExternalCompletion{}, err + } + published, err := h.completeExternalTool( + ctx, + m, + call, + outcome, + []Content{{Kind: "structured_data", Data: result}}, + ) + if err != nil { + return ExternalCompletion{}, err + } + if advance != nil && published.Matches { + _, err = dbsqlc.New(h.unit.DB()). + AdvanceProcessDefaultOutputCursor(ctx, + dbsqlc.AdvanceProcessDefaultOutputCursorParams{ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + ID: process.ID, + NextCursor: *advance}) + } + return published, err + }) +} + +func (h *Handle) replayProcessRead( + ctx context.Context, + call executiondb.ReadExecutionToolRow, + raw json.RawMessage, +) (ExternalCompletion, error) { + result := ExternalCompletion{ + Tool: CompletedTool{ID: call.ID, ResultID: valueOrZero(call.ResultID)}, + Published: true, + } + if valueOrZero(call.Outcome) != "succeeded" { + return result, nil + } + rows, err := executiondb.New(). + ReadExecutionContent(ctx, + h.unit.DB(), + executiondb.ReadExecutionContentParams{AgentID: h.route.AgentID, + ResultID: call.ResultID}) + if err != nil { + return result, err + } + for _, row := range rows { + if row.BlockKind != "structured_data" || row.StructuredData == nil { + continue + } + result.Matches, err = processresult.ReadMatches(raw, *row.StructuredData) + return result, err + } + return result, nil +} + +func (h *Handle) completeExternalTool( + ctx context.Context, + m *executionMutation, + call executiondb.ReadExecutionToolRow, + outcome string, + parts []Content, +) (ExternalCompletion, error) { + if call.Type != "built_in" || call.State != "waiting" && call.State != "completed" { + return ExternalCompletion{}, storeerr.ErrInvalidToolCallDisposition + } + completed, err := h.completeTool(ctx, m, call, outcome, parts, uuid.Nil, nil) + if errors.Is(err, storeerr.ErrIdempotencyConflict) { + return ExternalCompletion{ + Tool: CompletedTool{ID: call.ID, ResultID: valueOrZero(call.ResultID)}, + Published: call.State == "completed", + }, nil + } + return ExternalCompletion{Tool: completed, Matches: err == nil, Published: err == nil}, err +} + +func (h *Handle) FailProcessTool( + ctx context.Context, + processID uuid.UUID, + content []Content, +) (CompletedTool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (CompletedTool, error) { + process, err := executiondb.New(). + ReadExecutionProcessOwner(ctx, + h.unit.DB(), + executiondb.ReadExecutionProcessOwnerParams{AgentID: h.route.AgentID, + ID: processID}) + if err != nil { + return CompletedTool{}, err + } + call, err := h.tool(ctx, process.ToolCallID) + if err != nil { + return CompletedTool{}, err + } + if call.State != "waiting" || call.Type != "built_in" { + return CompletedTool{ID: call.ID}, nil + } + result, err := h.completeExternalTool(ctx, m, call, "failed", content) + return result.Tool, err + }) +} diff --git a/internal/storage/internal/agentexecution/commands_processes_integration_test.go b/internal/storage/internal/agentexecution/commands_processes_integration_test.go new file mode 100644 index 000000000..00ce4b473 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_processes_integration_test.go @@ -0,0 +1,374 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/processcmd" + "github.com/omnara-ai/omnara/internal/publicid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +func processOwner( + t *testing.T, + f executionFixture, + u *agentexecution.Unit, + agent, lease, tool uuid.UUID, + state string, +) uuid.UUID { + t.Helper() + machine, binding, id := uuid.New(), uuid.New(), uuid.New() + _, err := u.DB(). + Exec(t.Context(), + `INSERT INTO machines(id,org_id,display_name,source_kind,provider,lifecycle_state,lifecycle_changed_at,created_at,updated_at) + VALUES($1,$2,'test','byo','byo','active',statement_timestamp(),statement_timestamp(),statement_timestamp())`, + machine, f.org) + require.NoError(t, err) + _, err = u.DB(). + Exec(t.Context(), + `INSERT INTO agent_machine_bindings(id,org_id,project_id,agent_id,machine_id,machine_ref,binding_kind,created_at,updated_at) + VALUES($1,$2,$3,$4,$5,'mchr-abcdef','explicit',statement_timestamp(),statement_timestamp())`, + binding, f.org, f.project, agent, machine) + require.NoError(t, err) + spec, err := json.Marshal(processcmd.ForShell("echo done", "", "")) + require.NoError(t, err) + _, err = u.DB(). + Exec(t.Context(), + `INSERT INTO processes(id,org_id,project_id,agent_id,tool_call_id,runtime_lock_id,agent_machine_binding_id,machine_id, + execution_spec,state,execution_granted_at,source_started_at,state_changed_at,created_at,updated_at,last_activity_at) + VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,CASE WHEN $10='queued' THEN NULL ELSE statement_timestamp() END, + CASE WHEN $10='running' THEN statement_timestamp() ELSE NULL END, + statement_timestamp(),statement_timestamp(),statement_timestamp(),statement_timestamp())`, + id, f.org, f.project, agent, tool, lease, binding, machine, spec, state) + require.NoError(t, err) + return id +} + +func TestExecutionDurableProcessCustody(t *testing.T) { + for _, mode := range []string{"started", "terminal", "release", "reap", "cancel", "archive", "revoke"} { + t.Run(mode, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + ref := toolCommand(t, h, lease, "built_in") + _, err := h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + state := "queued" + if mode == "started" || mode == "terminal" { + state = "running" + } + process := processOwner(t, f, u, a, lease, ref.ID, state) + changed, err := h.WaitForTool(t.Context(), ref, agentexecution.ToolOwner{ProcessID: process}) + require.NoError(t, err) + require.True(t, changed) + changed, err = h.WaitForTool(t.Context(), ref, agentexecution.ToolOwner{ProcessID: process}) + require.NoError(t, err) + require.False(t, changed) + assertCommandState(t, u, h) + switch mode { + case "started": + input := agentexecution.ProcessResult{ + ID: process, + Started: true, + Observed: json.RawMessage(`{"output":"start","cursor":0,"next_cursor":5}`), + } + published, err := h.CompleteProcess(t.Context(), input) + require.NoError(t, err) + require.True(t, published.Matches) + require.True(t, published.Tool.Changed) + replay, err := h.CompleteProcess(t.Context(), input) + require.NoError(t, err) + require.True(t, replay.Matches) + require.False(t, replay.Tool.Changed) + _, err = u.DB(). + Exec(t.Context(), + `UPDATE processes SET state='exited',source_ended_at=statement_timestamp(),exit_code=0 WHERE id=$1`, + process) + require.NoError(t, err) + late, err := h.CompleteProcess(t.Context(), agentexecution.ProcessResult{ID: process}) + require.NoError(t, err) + require.False(t, late.Matches) + require.True(t, late.Published) + case "terminal": + work, err := h.Advance(t.Context(), lease, true, true) + require.NoError(t, err) + require.True(t, work.Released) + _, err = u.DB(). + Exec(t.Context(), + `UPDATE processes SET state='exited',source_ended_at=statement_timestamp(),exit_code=0 WHERE id=$1`, + process) + require.NoError(t, err) + published, err := h.CompleteProcess(t.Context(), agentexecution.ProcessResult{ID: process}) + require.NoError(t, err) + require.True(t, published.Matches) + case "release": + require.NoError(t, h.Release(t.Context(), lease)) + var state string + require.NoError( + t, + u.DB(). + QueryRow(t.Context(), `SELECT state FROM processes WHERE id=$1`, process). + Scan(&state), + ) + require.Equal(t, "queued", state) + case "reap": + _, err = u.DB(). + Exec(t.Context(), + `UPDATE agent_runtime_locks SET started_at=statement_timestamp()-interval '3 minutes',renewed_at=statement_timestamp()-interval '2 minutes',lease_expires_at=statement_timestamp()-interval '1 minute' WHERE id=$1`, + lease) + require.NoError(t, err) + reaped, err := h.Reap(t.Context(), lease) + require.NoError(t, err) + require.True(t, reaped) + case "cancel": + _, err = h.Cancel( + t.Context(), + agentexecution.CancelInput{Reason: "agent_canceled", Message: "canceled"}, + ) + require.NoError(t, err) + case "archive": + archived, err := h.Archive(t.Context(), uuid.Nil) + require.NoError(t, err) + require.True(t, archived) + case "revoke": + route, err := h.Route() + require.NoError(t, err) + require.NoError( + t, + u.InterruptProcesses( + t.Context(), + map[agentexecution.AgentRoute][]uuid.UUID{route: {process}}, + "grant_revoked", + ), + ) + } + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "process custody")) + }) + } +} + +func TestExecutionRejectsMissingDurableOwner(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + ref := toolCommand(t, h, lease, "built_in") + _, err := h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.WaitForTool(t.Context(), ref, agentexecution.ToolOwner{ProcessID: uuid.New()}) + return err + }), storeerr.ErrInvalidToolCallDisposition) + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.WorkTool, snapshot.Selection.Work) + require.NoError(t, u.Commit(t.Context(), "missing owner rollback")) +} + +func TestExecutionActionReports(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "actions") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = nil + var bindings []agentexecution.ToolProposal + for _, id := range []string{"start", "one", "two"} { + response.Normalized.Content = append( + response.Normalized.Content, + modelenvelope.ResponsePart{ + Type: "tool_call", + ProviderCallID: id, + ToolName: "tool", + ToolInput: json.RawMessage(`{}`), + }, + ) + bindings = append(bindings, agentexecution.ToolProposal{ProviderCallID: id, Type: "built_in"}) + } + output, err := h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ + RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: response, + Tools: bindings, + }, + ) + require.NoError(t, err) + refs := make([]agentexecution.ToolRef, 3) + for i, tool := range output.Tools { + refs[i] = agentexecution.ToolRef{ID: tool.ID, RuntimeLockID: lease} + _, err = h.AuthorizeTool(t.Context(), refs[i]) + require.NoError(t, err) + } + process := processOwner(t, f, u, a, lease, refs[0].ID, "running") + _, err = h.WaitForTool(t.Context(), refs[0], agentexecution.ToolOwner{ProcessID: process}) + require.NoError(t, err) + _, err = h.CompleteProcess(t.Context(), agentexecution.ProcessResult{ID: process, Started: true}) + require.NoError(t, err) + actions := []uuid.UUID{uuid.New(), uuid.New()} + for i, id := range actions { + _, err = u.DB(). + Exec(t.Context(), + `INSERT INTO process_actions(id,org_id,project_id,agent_id,process_id,tool_call_id,runtime_lock_id,action_kind,seq,payload,state,created_at,updated_at) + VALUES($1,$2,$3,$4,$5,$6,$7,'read',$8,'{}','queued',statement_timestamp(),statement_timestamp())`, + id, f.org, f.project, a, process, refs[i+1].ID, lease, i+1) + require.NoError(t, err) + _, err = h.WaitForTool(t.Context(), refs[i+1], agentexecution.ToolOwner{ActionID: id}) + require.NoError(t, err) + } + _, err = u.DB().Exec(t.Context(), `UPDATE process_actions SET state='applied' WHERE id=$1`, actions[1]) + require.NoError(t, err) + pid, err := publicid.Encode(publicid.KindProcess, process) + require.NoError(t, err) + observed := json.RawMessage( + `{"process_id":"` + pid + `","output":"ab","cursor":0,"next_cursor":2,"truncated":false}`, + ) + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + _, err := h.CompleteProcessAction( + t.Context(), + agentexecution.ActionResult{ID: actions[1], Observed: observed}, + ) + return err + }), storeerr.ErrProcessActionReportBlocked) + _, err = u.DB().Exec(t.Context(), `UPDATE process_actions SET state='applied' WHERE id=$1`, actions[0]) + require.NoError(t, err) + first, err := h.CompleteProcessAction( + t.Context(), + agentexecution.ActionResult{ID: actions[0], Observed: observed}, + ) + require.NoError(t, err) + require.True(t, first.Tool.Changed) + assertCommandState(t, u, h) + var cursor int64 + require.NoError( + t, + u.DB(). + QueryRow(t.Context(), `SELECT default_output_cursor FROM processes WHERE id=$1`, process). + Scan(&cursor), + ) + require.EqualValues(t, 2, cursor) + observedTwo := json.RawMessage( + `{"process_id":"` + pid + `","output":"cd","cursor":2,"next_cursor":4,"truncated":false}`, + ) + second, err := h.CompleteProcessAction( + t.Context(), + agentexecution.ActionResult{ID: actions[1], Observed: observedTwo}, + ) + require.NoError(t, err) + require.True(t, second.Matches) + _, err = u.DB(). + Exec(t.Context(), + `UPDATE processes SET state='exited',source_ended_at=statement_timestamp(),exit_code=0 WHERE id=$1`, + process) + require.NoError(t, err) + replay, err := h.CompleteProcessAction( + t.Context(), + agentexecution.ActionResult{ID: actions[0], Observed: observed}, + ) + require.NoError(t, err) + require.True(t, replay.Matches) + require.False(t, replay.Tool.Changed) + conflict, err := h.CompleteProcessAction( + t.Context(), + agentexecution.ActionResult{ID: actions[0], Observed: observedTwo}, + ) + require.NoError(t, err) + require.False(t, conflict.Matches) + assertCommandState(t, u, h) + + for i, state := range []string{"failed", "unknown"} { + prepared = prepareCommand(t, h, lease) + response.Normalized.Content = response.Normalized.Content[:1] + failedOutput, + err := h.AcceptOutput(t.Context(), + agentexecution.AcceptOutputInput{RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: response, + Tools: bindings[:1]}) + require.NoError(t, err) + ref := agentexecution.ToolRef{ID: failedOutput.Tools[0].ID, RuntimeLockID: lease} + _, err = h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + id := uuid.New() + kind := "read" + if state == "unknown" { + kind = "write" + } + _, err = u.DB(). + Exec(t.Context(), `INSERT INTO process_actions(id,org_id,project_id,agent_id,process_id,tool_call_id, + runtime_lock_id,action_kind,seq,payload,state,created_at,updated_at) + VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,'{}','queued',statement_timestamp(),statement_timestamp())`, + id, f.org, f.project, a, process, ref.ID, lease, kind, i+3) + require.NoError(t, err) + _, err = h.WaitForTool(t.Context(), ref, agentexecution.ToolOwner{ActionID: id}) + require.NoError(t, err) + _, err = u.DB().Exec(t.Context(), `UPDATE process_actions SET state=$2,state_reason_code='lost', + state_reason_message='connection lost' WHERE id=$1`, id, state) + require.NoError(t, err) + completion, err := h.CompleteProcessAction(t.Context(), agentexecution.ActionResult{ID: id}) + require.NoError(t, err) + require.True(t, completion.Matches) + require.True(t, completion.Tool.Changed) + repeated, err := h.CompleteProcessAction(t.Context(), agentexecution.ActionResult{ID: id}) + require.NoError(t, err) + require.True(t, repeated.Matches) + require.False(t, repeated.Tool.Changed) + var outcome string + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT outcome FROM tool_call_results WHERE id=$1`, + completion.Tool.ResultID).Scan(&outcome)) + require.Equal(t, "failed", outcome) + assertCommandState(t, u, h) + } + require.NoError(t, u.Commit(t.Context(), "ordered actions")) +} + +func TestExecutionOldTurnLateToolResult(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "old turn") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + old := prepareCommand(t, h, lease) + receiveCommand(t, h, "steering", "new turn") + admitted, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + require.NotEqual(t, old.Context.TurnID, admitted.TurnID) + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = []modelenvelope.ResponsePart{ + {Type: "tool_call", ProviderCallID: "late", ToolName: "tool", ToolInput: json.RawMessage(`{}`)}, + } + output, err := h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ + RuntimeLockID: lease, + ContextID: old.Context.ID, + Response: response, + Tools: []agentexecution.ToolProposal{{ProviderCallID: "late", Type: "custom"}}, + }, + ) + require.NoError(t, err) + ref := agentexecution.ToolRef{ID: output.Tools[0].ID, RuntimeLockID: lease} + _, err = h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + before := assertCommandState(t, u, h) + completed, err := h.CompleteCustomTool( + t.Context(), + ref.ID, + "succeeded", + []agentexecution.Content{{Kind: "text", Text: "late"}}, + ) + require.NoError(t, err) + require.Equal(t, old.Context.TurnID, completed.Event.TurnID) + after := assertCommandState(t, u, h) + require.Equal(t, before.Head, after.Head) + require.Equal(t, before.Selection, after.Selection) + require.NoError(t, u.Commit(t.Context(), "late result")) +} diff --git a/internal/storage/internal/agentexecution/commands_runtime.go b/internal/storage/internal/agentexecution/commands_runtime.go new file mode 100644 index 000000000..e2ef206f7 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_runtime.go @@ -0,0 +1,422 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type OwnedWork struct { + RuntimeID uuid.UUID + Selection Selection + Model PreparedModel + Admission InputAdmission + Released bool +} + +func (h *Handle) Claim( + ctx context.Context, + workerID uuid.UUID, + leaseDuration time.Duration, +) (OwnedWork, error) { + return executeCommand(ctx, h, func(m *executionMutation) (OwnedWork, error) { + if workerID == uuid.Nil || leaseDuration < 15*time.Second || leaseDuration > 30*time.Minute { + return OwnedWork{}, errors.New("worker identity is required") + } + q := executiondb.New() + _, err := q.ReadExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.ReadExecutionRuntimeParams{AgentID: h.route.AgentID}, + ) + if err == nil { + return OwnedWork{}, storeerr.ErrAgentNotAdvanceable + } + if !errors.Is(err, pgx.ErrNoRows) { + return OwnedWork{}, err + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return OwnedWork{}, err + } + if snapshot.Selection.Work == WorkNone { + m.changed() + return OwnedWork{Selection: snapshot.Selection}, nil + } + result := OwnedWork{} + if snapshot.Selection.Work == WorkInput { + result.Admission, err = h.admitInputs(ctx, m, snapshot.Selection.Admission) + if err != nil { + return OwnedWork{}, err + } + snapshot, err = h.LoadExecution(ctx) + if err != nil { + return OwnedWork{}, err + } + } + result.Selection = snapshot.Selection + if result.Selection.Work == WorkNone { + return result, nil + } + result.RuntimeID, err = q.AcquireExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.AcquireExecutionRuntimeParams{ + AgentID: h.route.AgentID, + ID: uuid.New(), + WorkerID: workerID, LeaseMicroseconds: leaseDuration.Microseconds(), + }, + ) + if err != nil { + return OwnedWork{}, err + } + h.unit.runtimeLocked = true + m.changed() + if err = q.ConsumeExecutionWakeup(ctx, + h.unit.DB(), + executiondb.ConsumeExecutionWakeupParams{AgentID: h.route.AgentID}); err != nil { + return OwnedWork{}, err + } + return result, nil + }) +} + +func (h *Handle) Advance( + ctx context.Context, + runtimeID uuid.UUID, + allowModelWork, prepareModel bool, +) (OwnedWork, error) { + return executeCommand(ctx, h, func(m *executionMutation) (OwnedWork, error) { + runtime, err := h.lockRuntime(ctx, runtimeID, false) + if err != nil { + return OwnedWork{}, err + } + if runtime.CancelRequestedAt != nil || !runtime.LeaseExpiresAt.After(runtime.DatabaseNow) { + return OwnedWork{}, storeerr.ErrRuntimeLockInactive + } + retained, err := executiondb.New(). + ExecutionRetainedRuntimeWork(ctx, + h.unit.DB(), + executiondb.ExecutionRetainedRuntimeWorkParams{AgentID: h.route.AgentID, + RuntimeID: &runtimeID}) + if err != nil { + return OwnedWork{}, err + } + if valueOrZero(retained) { + if err = h.releaseRuntime(ctx, m, runtime, false); err != nil { + return OwnedWork{}, err + } + return OwnedWork{Released: true}, nil + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return OwnedWork{}, err + } + result := OwnedWork{RuntimeID: runtimeID} + if !allowModelWork && (snapshot.Selection.Work == WorkModel || snapshot.Selection.Work == WorkInput) { + if err = h.releaseRuntime(ctx, m, runtime, false); err != nil { + return OwnedWork{}, err + } + return OwnedWork{Released: true}, nil + } + if snapshot.Selection.Work == WorkInput { + result.Admission, err = h.admitInputs(ctx, m, snapshot.Selection.Admission) + if err != nil { + return OwnedWork{}, err + } + snapshot, err = h.LoadExecution(ctx) + if err != nil { + return OwnedWork{}, err + } + } + result.Selection = snapshot.Selection + if snapshot.Selection.Work == WorkNone { + if err = h.releaseRuntime(ctx, m, runtime, false); err != nil { + return OwnedWork{}, err + } + result.RuntimeID = uuid.Nil + result.Released = true + return result, nil + } + if prepareModel && snapshot.Selection.Work == WorkModel && + snapshot.Selection.Model.Kind != ModelResume { + result.Model, err = h.PrepareModel( + ctx, + PrepareModelInput{RuntimeLockID: runtimeID, Selected: *snapshot.Selection.Model}, + ) + if err != nil { + return OwnedWork{}, err + } + } + m.changed() + return result, nil + }) +} + +func (h *Handle) AcceptAndAdvance( + ctx context.Context, + input AcceptOutputInput, + allowModelWork, prepareModel bool, +) (AcceptedOutput, OwnedWork, error) { + output, err := h.AcceptOutput(ctx, input) + if err != nil { + return AcceptedOutput{}, OwnedWork{}, err + } + if !output.Created { + return output, OwnedWork{Released: true}, h.Release(ctx, input.RuntimeLockID) + } + work, err := h.Advance(ctx, input.RuntimeLockID, allowModelWork, prepareModel) + return output, work, err +} + +func (h *Handle) lockRuntime( + ctx context.Context, + id uuid.UUID, + skip bool, +) (executiondb.ReadExecutionRuntimeRow, error) { + q := executiondb.New() + var err error + if skip { + _, err = q.TryLockExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.TryLockExecutionRuntimeParams{AgentID: h.route.AgentID, ID: id}, + ) + } else { + _, + err = q.LockExecutionRuntime(ctx, + h.unit.DB(), + executiondb.LockExecutionRuntimeParams{AgentID: h.route.AgentID, + ID: id}) + } + if errors.Is(err, pgx.ErrNoRows) { + return executiondb.ReadExecutionRuntimeRow{}, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return executiondb.ReadExecutionRuntimeRow{}, err + } + h.unit.runtimeLocked = true + return q.ReadExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.ReadExecutionRuntimeParams{AgentID: h.route.AgentID, ID: &id}, + ) +} + +func (h *Handle) Release(ctx context.Context, id uuid.UUID) error { + _, err := executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + runtime, err := h.lockRuntime(ctx, id, false) + if err != nil { + return false, err + } + return true, h.releaseRuntime(ctx, m, runtime, false) + }) + return err +} + +func (h *Handle) Reap(ctx context.Context, id uuid.UUID) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + runtime, err := h.lockRuntime(ctx, id, true) + if errors.Is(err, storeerr.ErrRuntimeLockInactive) { + return false, nil + } + if err != nil { + return false, err + } + if runtime.LeaseExpiresAt.After(runtime.DatabaseNow) { + return false, nil + } + return true, h.releaseRuntime(ctx, m, runtime, true) + }) +} + +type RuntimeRenewal struct { + LeaseExpiresAt time.Time + LocalStartedAt time.Time +} + +func (u *Unit) RenewRuntime( + ctx context.Context, route AgentRoute, id uuid.UUID, leaseDuration time.Duration, +) (RuntimeRenewal, error) { + if route.CellID != u.cell.id || route.ProjectID == uuid.Nil || route.AgentID == uuid.Nil || + id == uuid.Nil { + return RuntimeRenewal{}, storeerr.ErrRuntimeLockInactive + } + if leaseDuration < 15*time.Second || leaseDuration > 30*time.Minute { + return RuntimeRenewal{}, errors.New("runtime lease must be between 15 seconds and 30 minutes") + } + q := executiondb.New() + _, err := q.LockExecutionRuntime( + ctx, + u.DB(), + executiondb.LockExecutionRuntimeParams{AgentID: route.AgentID, ID: id}, + ) + if errors.Is(err, pgx.ErrNoRows) { + return RuntimeRenewal{}, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return RuntimeRenewal{}, err + } + u.runtimeLocked = true + started := time.Now() //nolint:omnaralint // Monotonic lease budget. + expires, err := q.RenewExecutionRuntime(ctx, u.DB(), executiondb.RenewExecutionRuntimeParams{ + ProjectID: route.ProjectID, AgentID: route.AgentID, ID: id, + LeaseMicroseconds: leaseDuration.Microseconds(), + }) + if errors.Is(err, pgx.ErrNoRows) { + return RuntimeRenewal{}, storeerr.ErrRuntimeLockInactive + } + return RuntimeRenewal{LeaseExpiresAt: expires, LocalStartedAt: started}, err +} + +func (h *Handle) releaseRuntime( + ctx context.Context, + m *executionMutation, + runtime executiondb.ReadExecutionRuntimeRow, + expired bool, +) error { + code := "runtime_released_before_model_result_acceptance" + message := "runtime released before the model result was durably accepted" + reason := "runtime_lock_released" + if expired { + code = "runtime_lease_expired_before_model_result_acceptance" + message = "runtime lease expired before the model result was durably accepted" + reason = "runtime_lock_stale" + } + q := executiondb.New() + live, err := q.ListExecutionLiveContexts( + ctx, + h.unit.DB(), + executiondb.ListExecutionLiveContextsParams{AgentID: h.route.AgentID, RuntimeID: &runtime.ID}, + ) + if err != nil { + return err + } + for _, c := range live { + row, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: c.ID}, + ) + if err != nil { + return err + } + details, err := json.Marshal( + map[string]any{ + "attempt_number": row.AttemptNumber, + "code": code, + "message": message, + "outcome_ambiguous": true, + "runtime_lock_id": runtime.ID, + }, + ) + if err != nil { + return err + } + args := executiondb.TeardownExecutionContextParams{ + AgentID: h.route.AgentID, + ID: c.ID, + RuntimeID: runtime.ID, + State: "failed", + ErrorKind: "runtime", + ErrorCode: code, + ErrorMessage: message, + ErrorDetails: details, + } + if row.AttemptNumber <= 8 { + delay := max(0, h.unit.cell.retryBackoff(int(row.AttemptNumber), c.ID.String()).Microseconds()) + args.RetryMicroseconds = &delay + args.Recovery = optionalText("retry") + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return err + } + n, err := q.TeardownExecutionContext(ctx, h.unit.DB(), args) + if err != nil { + return err + } + if n != 1 { + return storeerr.ErrStateTransitionConflict + } + m.changed() + if args.Recovery == nil { + _, err = h.publishTerminalFailure( + ctx, + m, + ModelFailure{ + ContextID: c.ID, + RuntimeLockID: runtime.ID, + ErrorKind: "runtime", + ErrorCode: code, + ErrorMessage: message, + ErrorDetails: details, + }, + row, + snapshot, + true, + ) + if err != nil { + return err + } + } + } + calls, err := q.ListExecutionUnfinishedTools( + ctx, + h.unit.DB(), + executiondb.ListExecutionUnfinishedToolsParams{AgentID: h.route.AgentID, RuntimeID: &runtime.ID}, + ) + if err != nil { + return err + } + body, err := json.Marshal( + map[string]string{ + "error_code": reason, + "message": "Tool call was interrupted; its external outcome is unknown.", + }, + ) + if err != nil { + return err + } + for _, id := range calls { + call, err := h.tool(ctx, id) + if err != nil { + return err + } + if _, + err = h.completeTool(ctx, + m, + call, + "failed", + []Content{{Kind: "structured_data", + Data: body}}, + uuid.Nil, nil); err != nil { + return err + } + } + if expired { + if err = h.failQueuedRuntime(ctx, runtime.ID, reason); err != nil { + return err + } + } + n, err := q.DeleteExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.DeleteExecutionRuntimeParams{AgentID: h.route.AgentID, ID: runtime.ID}, + ) + if err != nil { + return err + } + if n != 1 { + return storeerr.ErrRuntimeLockInactive + } + m.changed() + return nil +} diff --git a/internal/storage/internal/agentexecution/commands_runtime_integration_test.go b/internal/storage/internal/agentexecution/commands_runtime_integration_test.go new file mode 100644 index 000000000..e27b056aa --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_runtime_integration_test.go @@ -0,0 +1,228 @@ +//go:build integration + +package agentexecution_test + +import ( + "context" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +func TestExecutionAdvanceBudgetsAndRecovery(t *testing.T) { + for _, mode := range []string{"tools", + "model_budget", + "input_budget", + "retained", + "model", + "canceled", + "expired"} { + t.Run(mode, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + ref := toolCommand(t, h, lease, "built_in") + _, err := h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + if mode == "retained" { + _, err = h.RunTool(t.Context(), ref) + require.NoError(t, err) + } + if mode == "model_budget" || mode == "input_budget" || mode == "model" { + _, err = h.CompleteTool( + t.Context(), + agentexecution.ToolCompletion{ + ToolRef: ref, + Outcome: "succeeded", + Content: []agentexecution.Content{{Kind: "text", Text: "done"}}, + }, + ) + require.NoError(t, err) + } + if mode == "input_budget" { + receiveCommand(t, h, "steering", "new input") + } + if mode == "canceled" || mode == "expired" { + query := `UPDATE agent_runtime_locks SET cancel_requested_at=statement_timestamp() WHERE id=$1` + if mode == "expired" { + query = `UPDATE agent_runtime_locks SET started_at=statement_timestamp()-interval '3 minutes', + renewed_at=statement_timestamp()-interval '2 minutes', + lease_expires_at=statement_timestamp()-interval '1 minute' WHERE id=$1` + } + _, err = u.DB().Exec(t.Context(), query, lease) + require.NoError(t, err) + require.ErrorIs( + t, + u.Savepoint( + t.Context(), + func(_ *agentexecution.Unit) error { _, err := h.Advance(t.Context(), lease, true, true); return err }, + ), + storeerr.ErrRuntimeLockInactive, + ) + require.NoError(t, u.Commit(t.Context(), "fenced advance")) + return + } + work, err := h.Advance(t.Context(), lease, mode == "model", true) + require.NoError(t, err) + switch mode { + case "tools": + require.Equal(t, agentexecution.WorkTool, work.Selection.Work) + require.False(t, work.Released) + case "model": + require.Equal(t, agentexecution.WorkModel, work.Selection.Work) + require.NotEqual(t, uuid.Nil, work.Model.Context.ID) + default: + require.True(t, work.Released) + } + snapshot := assertCommandState(t, u, h) + if mode == "input_budget" { + require.Equal(t, agentexecution.AdmitAllSteering, snapshot.Selection.Admission) + } + if mode == "retained" { + require.Equal(t, agentexecution.WorkModel, snapshot.Selection.Work) + } + require.NoError(t, u.Commit(t.Context(), "advance")) + }) + } +} + +func TestExecutionClaimRenewAcceptAdvance(t *testing.T) { + f, a, oldLease := commandFixture(t) + _, err := f.pool.Exec(t.Context(), `DELETE FROM agent_runtime_locks WHERE id=$1`, oldLease) + require.NoError(t, err) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "claim") + claimed, err := h.Claim(t.Context(), uuid.New(), 90*time.Second) + require.NoError(t, err) + require.NotEqual(t, uuid.Nil, claimed.RuntimeID) + require.ErrorIs( + t, + u.Savepoint( + t.Context(), + func(_ *agentexecution.Unit) error { + _, err := h.Claim(t.Context(), uuid.New(), 90*time.Second) + return err + }, + ), + storeerr.ErrAgentNotAdvanceable, + ) + route, err := h.Route() + require.NoError(t, err) + _, err = u.RenewRuntime(t.Context(), route, claimed.RuntimeID, 90*time.Second) + require.NoError(t, err) + prepared := prepareCommand(t, h, claimed.RuntimeID) + input := agentexecution.AcceptOutputInput{ + RuntimeLockID: claimed.RuntimeID, + ContextID: prepared.Context.ID, + Response: commandResponse(modelenvelope.StopReasonMaxTokens), + } + accepted, work, err := h.AcceptAndAdvance(t.Context(), input, true, true) + require.NoError(t, err) + require.True(t, accepted.Created) + require.NotEqual(t, prepared.Context.ID, work.Model.Context.ID) + replay, work, err := h.AcceptAndAdvance(t.Context(), input, true, true) + require.NoError(t, err) + require.False(t, replay.Created) + require.Equal(t, uuid.Nil, work.Model.Context.ID) + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "claim renew advance")) +} + +func TestExecutionRecoveryExhaustion(t *testing.T) { + for _, compact := range []bool{false, true} { + t.Run(map[bool]string{false: "normal", true: "compaction"}[compact], func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "exhaustion") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + if compact { + start, + err := h.BeginCompaction(t.Context(), + agentexecution.BeginCompactionInput{SourceEnd: prepared.Context.InputEventSequence, + Failure: agentexecution.ModelFailure{ContextID: prepared.Context.ID, + RuntimeLockID: lease, + ErrorKind: "context_length", + ErrorMessage: "compact"}}) + require.NoError(t, err) + prepared = start.Model + } + for range 8 { + _, + err = h.FailModel(t.Context(), + agentexecution.ModelFailure{ContextID: prepared.Context.ID, + RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, + ErrorKind: "transient", + ErrorMessage: "retry"}) + require.NoError(t, err) + prepared = prepareCommand(t, h, lease) + } + require.EqualValues(t, 9, prepared.Context.Attempt) + require.NoError(t, h.Release(t.Context(), lease)) + snapshot := assertCommandState(t, u, h) + require.False(t, snapshot.Selection.TurnContinuable) + require.Equal(t, agentexecution.WorkNone, snapshot.Selection.Work) + var outputs int + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT count(*) FROM model_outputs WHERE agent_id=$1 AND stop_reason='error'`, + a).Scan(&outputs)) + require.Equal(t, 1, outputs) + reaped, err := h.Reap(t.Context(), lease) + require.NoError(t, err) + require.False(t, reaped) + require.NoError(t, u.Commit(t.Context(), "recovery exhausted")) + }) + } +} + +func TestExecutionRenewalSamplesTimeAfterLock(t *testing.T) { + f, a, lease := commandFixture(t) + _, err := f.pool.Exec(t.Context(), `UPDATE agent_runtime_locks SET + lease_expires_at=statement_timestamp()+interval '2 seconds' WHERE id=$1`, lease) + require.NoError(t, err) + blocker, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = blocker.Rollback(context.Background()) }() + _, err = blocker.Exec(t.Context(), `SELECT id FROM agent_runtime_locks WHERE id=$1 FOR UPDATE`, lease) + require.NoError(t, err) + cell := agentexecution.NewCell("test", f.pool, nil, nil) + result := make(chan error, 1) + go func() { + u, err := cell.Begin(t.Context()) + if err != nil { + result <- err + return + } + defer func() { _ = u.Rollback(context.Background()) }() + _, + err = u.RenewRuntime(t.Context(), + agentexecution.AgentRoute{CellID: "test", + ProjectID: f.project, + AgentID: a, + RootAgentID: a}, + lease, + 90*time.Second) + result <- err + }() + require.Eventually(t, func() bool { + var blocked bool + err := f.pool.QueryRow(t.Context(), `SELECT EXISTS(SELECT 1 FROM pg_stat_activity + WHERE datname=current_database() AND cardinality(pg_blocking_pids(pid))>0)`).Scan(&blocked) + return err == nil && blocked + }, time.Second, 10*time.Millisecond) + require.Eventually(t, func() bool { + var expired bool + err := f.pool.QueryRow(t.Context(), `SELECT lease_expires_at<=statement_timestamp() + FROM agent_runtime_locks WHERE id=$1`, lease).Scan(&expired) + return err == nil && expired + }, 3*time.Second, 10*time.Millisecond) + require.NoError(t, blocker.Commit(t.Context())) + require.NoError(t, <-result) +} diff --git a/internal/storage/internal/agentexecution/commands_selection_integration_test.go b/internal/storage/internal/agentexecution/commands_selection_integration_test.go new file mode 100644 index 000000000..60cc6a914 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_selection_integration_test.go @@ -0,0 +1,111 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func TestExecutionBatchConsumption(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "initial") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + for cycle := range 4 { + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = []modelenvelope.ResponsePart{ + {Type: "tool_call", ProviderCallID: "one", ToolName: "tool", ToolInput: json.RawMessage(`{}`)}, + {Type: "tool_call", ProviderCallID: "two", ToolName: "tool", ToolInput: json.RawMessage(`{}`)}, + } + output, err := h.AcceptOutput(t.Context(), agentexecution.AcceptOutputInput{ + ContextID: prepared.Context.ID, RuntimeLockID: lease, Response: response, + Tools: []agentexecution.ToolProposal{ + {ProviderCallID: "one", Type: "built_in"}, {ProviderCallID: "two", Type: "built_in"}, + }, + }) + require.NoError(t, err) + var config agentexecution.ActivatedConfig + if cycle == 1 { + config, err = h.ActivateConfig(t.Context(), agentexecution.ActivateConfigInput{ID: uuid.New(), ConfigID: f.config}) + require.NoError(t, err) + } + for _, tool := range output.Tools { + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.WorkTool, snapshot.Selection.Work) + require.Equal(t, output.ID, snapshot.Head.PendingToolOutputID) + ref := agentexecution.ToolRef{ID: tool.ID, RuntimeLockID: lease} + _, err = h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + _, err = h.CompleteTool(t.Context(), agentexecution.ToolCompletion{ToolRef: ref, Outcome: "succeeded", + Content: []agentexecution.Content{{Kind: "text", Text: "done"}}}) + require.NoError(t, err) + } + snapshot := assertCommandState(t, u, h) + if cycle == 1 { + require.Equal(t, agentexecution.OriginConfig, snapshot.Selection.Model.Origin) + require.Equal(t, config.InputID, snapshot.Selection.Model.SourceInputID) + } else { + require.Equal(t, agentexecution.OriginTools, snapshot.Selection.Model.Origin) + } + next := prepareCommand(t, h, lease) + require.GreaterOrEqual(t, next.Context.InputEventSequence, snapshot.View.ToolBatch.Completion.LastResultSequence) + snapshot = assertCommandState(t, u, h) + require.Equal(t, output.ID, snapshot.Head.PendingToolOutputID) + require.Equal(t, agentexecution.OriginTools, snapshot.Selection.Model.Origin) + _, err = h.FailModel(t.Context(), agentexecution.ModelFailure{ContextID: next.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, ErrorKind: "transient", ErrorMessage: "retry"}) + require.NoError(t, err) + snapshot = assertCommandState(t, u, h) + require.Equal(t, agentexecution.OriginRetry, snapshot.Selection.Model.Origin) + prepared = prepareCommand(t, h, lease) + snapshot = assertCommandState(t, u, h) + require.Equal(t, output.ID, snapshot.Head.PendingToolOutputID) + } + _, err = h.AcceptOutput(t.Context(), agentexecution.AcceptOutputInput{ContextID: prepared.Context.ID, + RuntimeLockID: lease, Response: commandResponse(modelenvelope.StopReasonEndTurn)}) + require.NoError(t, err) + snapshot := assertCommandState(t, u, h) + require.Nil(t, snapshot.View.ToolBatch) + require.Nil(t, snapshot.Selection.Model) + require.NoError(t, u.Commit(t.Context(), "batch consumption")) +} + +func TestExecutionSteeringAcrossPreparationAndRetry(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + initial := receiveCommand(t, h, "queued", "initial") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + first := receiveCommand(t, h, "steering", "after claim") + prepared := prepareCommand(t, h, lease) + require.Equal(t, []uuid.UUID{initial.ID}, prepared.Context.Opening.InputIDs) + _, err = h.FailModel(t.Context(), agentexecution.ModelFailure{ContextID: prepared.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, ErrorKind: "transient", ErrorMessage: "retry"}) + require.NoError(t, err) + retry := prepareCommand(t, h, lease) + require.Equal(t, prepared.Context.Opening, retry.Context.Opening) + require.Equal(t, prepared.Context.InputEventSequence, retry.Context.InputEventSequence) + _, err = h.FailModel(t.Context(), agentexecution.ModelFailure{ContextID: retry.Context.ID, RuntimeLockID: lease, + Recovery: agentexecution.RecoveryRetry, RetryDelay: time.Hour, ErrorKind: "transient", ErrorMessage: "retry later"}) + require.NoError(t, err) + receiveCommand(t, h, "queued", "later queued") + second := receiveCommand(t, h, "steering", "second steering") + snapshot := assertCommandState(t, u, h) + require.Equal(t, agentexecution.AdmitAllSteering, snapshot.Selection.Admission) + require.Equal(t, agentexecution.OriginRetry, snapshot.Selection.Model.Origin) + admitted, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + require.Len(t, admitted.Inputs, 2) + next := prepareCommand(t, h, lease) + require.Equal(t, []uuid.UUID{initial.ID, first.ID, second.ID}, next.Context.Opening.InputIDs) + require.NoError(t, u.Commit(t.Context(), "steering boundary")) +} diff --git a/internal/storage/internal/agentexecution/commands_tools.go b/internal/storage/internal/agentexecution/commands_tools.go new file mode 100644 index 000000000..ddc66c0a2 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_tools.go @@ -0,0 +1,339 @@ +package agentexecution + +import ( + "context" + "errors" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type ToolRef struct { + ID uuid.UUID + RuntimeLockID uuid.UUID +} +type ToolOwner struct { + InteractionID uuid.UUID + ProcessID uuid.UUID + ActionID uuid.UUID +} +type ToolCompletion struct { + ToolRef + Outcome string + Content []Content +} +type CompletedTool struct { + ID uuid.UUID + ResultID uuid.UUID + Event ExecutionEvent + Changed bool +} + +func (h *Handle) tool(ctx context.Context, id uuid.UUID) (executiondb.ReadExecutionToolRow, error) { + return executiondb.New(). + ReadExecutionTool(ctx, h.unit.DB(), executiondb.ReadExecutionToolParams{AgentID: h.route.AgentID, ID: id}) +} + +func (h *Handle) transitionTool( + ctx context.Context, + m *executionMutation, + row executiondb.ReadExecutionToolRow, + state string, + owner, fence uuid.UUID, +) error { + n, err := executiondb.New(). + TransitionExecutionTool(ctx, h.unit.DB(), executiondb.TransitionExecutionToolParams{ + AgentID: h.route.AgentID, ID: row.ID, PreviousState: row.State, PreviousRuntimeID: row.RuntimeLockID, + NextState: state, NextRuntimeID: nullableID(owner), FenceID: nullableID(fence)}) + if err != nil { + return err + } + if n != 1 { + return storeerr.ErrStateTransitionConflict + } + m.changed() + return nil +} + +func (h *Handle) AuthorizeTool(ctx context.Context, ref ToolRef) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + if err := h.FenceRuntime(ctx, ref.RuntimeLockID); err != nil { + return false, err + } + row, err := h.tool(ctx, ref.ID) + if err != nil { + return false, err + } + if row.State == "ready" { + return false, nil + } + if row.State != "awaiting_authorization" { + return false, storeerr.ErrStateTransitionConflict + } + if err = h.transitionTool(ctx, m, row, "ready", uuid.Nil, ref.RuntimeLockID); err != nil { + return false, err + } + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, row.ID, "ready", nil) + return true, nil + }) +} + +func (h *Handle) RunTool(ctx context.Context, ref ToolRef) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + if err := h.FenceRuntime(ctx, ref.RuntimeLockID); err != nil { + return false, err + } + row, err := h.tool(ctx, ref.ID) + if err != nil { + return false, err + } + if row.State == "running" { + return false, storeerr.ErrToolCallInProgress + } + if row.State != "ready" || row.Type == "custom" { + return false, storeerr.ErrStateTransitionConflict + } + if err = h.transitionTool(ctx, m, row, "running", ref.RuntimeLockID, ref.RuntimeLockID); err != nil { + return false, err + } + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, row.ID, "running", nil) + return true, nil + }) +} + +func (h *Handle) RequeueTool(ctx context.Context, ref ToolRef) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + if err := h.FenceRuntime(ctx, ref.RuntimeLockID); err != nil { + return false, err + } + row, err := h.tool(ctx, ref.ID) + if err != nil { + return false, err + } + if row.State == "ready" { + return false, nil + } + if row.State != "running" || valueOrZero(row.RuntimeLockID) != ref.RuntimeLockID { + return false, storeerr.ErrStateTransitionConflict + } + if err = h.transitionTool(ctx, m, row, "ready", uuid.Nil, ref.RuntimeLockID); err != nil { + return false, err + } + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, row.ID, "ready", nil) + return true, nil + }) +} + +func (h *Handle) WaitForTool(ctx context.Context, ref ToolRef, owner ToolOwner) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + owners := 0 + for _, id := range []uuid.UUID{owner.InteractionID, owner.ProcessID, owner.ActionID} { + if id != uuid.Nil { + owners++ + } + } + if owners != 1 { + return false, storeerr.ErrInvalidToolCallDisposition + } + + if err := h.FenceRuntime(ctx, ref.RuntimeLockID); err != nil { + return false, err + } + row, err := h.tool(ctx, ref.ID) + if err != nil { + return false, err + } + durable, err := executiondb.New(). + ExecutionToolOwner(ctx, h.unit.DB(), executiondb.ExecutionToolOwnerParams{ + AgentID: h.route.AgentID, ToolID: row.ID, InteractionID: nullableID(owner.InteractionID), + ProcessID: nullableID(owner.ProcessID), ActionID: nullableID(owner.ActionID)}) + if err != nil { + return false, err + } + if row.Type != "built_in" || !durable.Question && !durable.Process && !durable.Action { + return false, storeerr.ErrInvalidToolCallDisposition + } + if row.State == "waiting" { + return false, nil + } + if row.State != "ready" || row.Type != "built_in" { + return false, storeerr.ErrStateTransitionConflict + } + if err = h.transitionTool(ctx, m, row, "waiting", uuid.Nil, ref.RuntimeLockID); err != nil { + return false, err + } + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, row.ID, "waiting", nil) + return true, nil + }) +} + +func (h *Handle) CompleteTool(ctx context.Context, input ToolCompletion) (CompletedTool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (CompletedTool, error) { + if err := h.FenceRuntime(ctx, input.RuntimeLockID); err != nil { + return CompletedTool{}, err + } + row, err := h.tool(ctx, input.ID) + if err != nil { + return CompletedTool{}, err + } + allowed := row.State == "completed" || (row.State == "ready" && row.Type != "custom") || + (row.State == "running" && valueOrZero(row.RuntimeLockID) == input.RuntimeLockID) || + ((row.State == "awaiting_authorization" || + row.State == "awaiting_permission") && + input.Outcome != "succeeded") + if !allowed { + return CompletedTool{}, storeerr.ErrStateTransitionConflict + } + return h.completeTool(ctx, m, row, input.Outcome, input.Content, input.RuntimeLockID, nil) + }) +} + +func (h *Handle) CompleteCustomTool( + ctx context.Context, + id uuid.UUID, + outcome string, + content []Content, +) (CompletedTool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (CompletedTool, error) { + row, err := h.tool(ctx, id) + if err != nil { + return CompletedTool{}, err + } + if row.Type != "custom" || (row.State != "ready" && row.State != "completed") || + (outcome != "succeeded" && outcome != "failed") { + return CompletedTool{}, storeerr.ErrStateTransitionConflict + } + return h.completeTool(ctx, m, row, outcome, content, uuid.Nil, nil) + }) +} + +func (h *Handle) completeTool( + ctx context.Context, + m *executionMutation, + row executiondb.ReadExecutionToolRow, + outcome string, + content []Content, + fence uuid.UUID, + interaction *notifications.AgentInteractionUpdate, +) (CompletedTool, error) { + switch outcome { + case "succeeded", "failed", "denied", "canceled": + default: + return CompletedTool{}, errors.New("invalid tool outcome") + } + for _, part := range content { + if part.Kind != "text" && part.Kind != "structured_data" && part.Kind != "artifact" { + return CompletedTool{}, errors.New("invalid tool result content") + } + } + if row.State == "completed" { + if valueOrZero(row.Outcome) != outcome || row.ResultID == nil || row.ResultEventID == nil { + return CompletedTool{}, storeerr.ErrIdempotencyConflict + } + err := h.matchContent( + ctx, + executiondb.ReadExecutionContentParams{AgentID: h.route.AgentID, ResultID: row.ResultID}, + content, + ) + return CompletedTool{ + ID: row.ID, + ResultID: *row.ResultID, + Event: ExecutionEvent{ID: *row.ResultEventID, TurnID: row.TurnID, + EventBoundary: EventBoundary{ + Sequence: valueOrZero(row.ResultSequence), + Time: valueOrZero(row.CompletedAt), + }}, + }, err + } + if err := h.transitionTool(ctx, m, row, "completed", uuid.Nil, fence); err != nil { + return CompletedTool{}, err + } + q := executiondb.New() + id, err := q.InsertExecutionToolResult(ctx, h.unit.DB(), executiondb.InsertExecutionToolResultParams{ + AgentID: h.route.AgentID, ID: row.ID, Outcome: outcome}) + if err != nil { + return CompletedTool{}, err + } + if err = h.writeContent(ctx, "tool_call_result", id, content); err != nil { + return CompletedTool{}, err + } + event, err := q.AppendExecutionEvent(ctx, h.unit.DB(), executiondb.AppendExecutionEventParams{ + ID: uuid.New(), + ProjectID: h.route.ProjectID, + AgentID: h.route.AgentID, + TurnID: &row.TurnID, + Kind: "tool_result", + ResultID: &id, + IdempotencyKey: optionalText("tool-call-result-authority:" + row.ID.String())}) + if err != nil { + return CompletedTool{}, err + } + admitted := ExecutionEvent{ + ID: event.ID, + TurnID: row.TurnID, + EventBoundary: EventBoundary{Sequence: event.Sequence, Time: event.CreatedAt}, + } + if err = h.advanceTurn(ctx, admitted, false); err != nil { + return CompletedTool{}, err + } + closed, err := q.CloseExecutionToolInteractions( + ctx, + h.unit.DB(), + executiondb.CloseExecutionToolInteractionsParams{ + AgentID: h.route.AgentID, ToolID: row.ID}, + ) + if err != nil { + return CompletedTool{}, err + } + for _, i := range closed { + interaction = ¬ifications.AgentInteractionUpdate{ + ID: i.ID, + InteractionKind: i.InteractionKind, + State: i.State, + } + } + h.unit.Notifications().AddAgentEvent(h.route.AgentID, event.Sequence, "tool_result") + h.unit.Notifications().AddToolCallUpdate(h.route.AgentID, row.ID, "completed", interaction) + return CompletedTool{ID: row.ID, ResultID: id, Event: admitted, Changed: true}, nil +} + +type ToolDispatch uint8 + +const ( + ToolDispatchReady ToolDispatch = iota + 1 + ToolDispatchWaiting + ToolDispatchCompleted +) + +func (h *Handle) DispatchTool(ctx context.Context, ref ToolRef, acceptExisting bool) (ToolDispatch, error) { + if err := h.unit.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: h.route.ProjectID, AgentID: h.route.AgentID}}, + RuntimeAuthority{AgentID: h.route.AgentID, RuntimeLockID: ref.RuntimeLockID}, + ); err != nil { + return 0, err + } + state, err := executiondb.New().ReadExecutionToolDispatch( + ctx, h.unit.DB(), executiondb.ReadExecutionToolDispatchParams{AgentID: h.route.AgentID, ID: ref.ID}, + ) + if err != nil { + return 0, err + } + switch state { + case "ready": + return ToolDispatchReady, nil + case "running": + return 0, storeerr.ErrToolCallInProgress + case "waiting": + if acceptExisting { + return ToolDispatchWaiting, nil + } + case "completed": + if acceptExisting { + return ToolDispatchCompleted, nil + } + } + return 0, storeerr.ErrIdempotencyConflict +} diff --git a/internal/storage/internal/agentexecution/database.go b/internal/storage/internal/agentexecution/database.go new file mode 100644 index 000000000..e421c5129 --- /dev/null +++ b/internal/storage/internal/agentexecution/database.go @@ -0,0 +1,138 @@ +package agentexecution + +import ( + "context" + "errors" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type unitDB struct{ unit *Unit } + +func (db unitDB) Exec(ctx context.Context, query string, args ...any) (pgconn.CommandTag, error) { + if err := db.unit.active(); err != nil { + return pgconn.CommandTag{}, err + } + tag, err := db.unit.tx.Exec(ctx, query, args...) + db.unit.recordError(err) + return tag, err +} + +func (db unitDB) Query(ctx context.Context, query string, args ...any) (pgx.Rows, error) { + if err := db.unit.active(); err != nil { + return nil, err + } + rows, err := db.unit.tx.Query(ctx, query, args...) + db.unit.recordError(err) + if err != nil { + return nil, err + } + return unitRows{Rows: rows, unit: db.unit}, nil +} + +func (db unitDB) QueryRow(ctx context.Context, query string, args ...any) pgx.Row { + if err := db.unit.active(); err != nil { + return errorRow{err: err} + } + return unitRow{row: db.unit.tx.QueryRow(ctx, query, args...), unit: db.unit} +} + +func (u *Unit) recordError(err error) { + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) || errors.Is(err, context.Canceled) || + errors.Is(err, context.DeadlineExceeded) { + u.aborted = true + } +} + +type errorRow struct{ err error } + +func (row errorRow) Scan(...any) error { return row.err } + +type unitRow struct { + row pgx.Row + unit *Unit +} + +func (row unitRow) Scan(dest ...any) error { + err := row.row.Scan(dest...) + row.unit.recordError(err) + return err +} + +type unitRows struct { + pgx.Rows + unit *Unit +} + +func (rows unitRows) Conn() *pgx.Conn { return nil } + +func (rows unitRows) Next() bool { + next := rows.Rows.Next() + if !next { + rows.unit.recordError(rows.Rows.Err()) + } + return next +} + +func (rows unitRows) Err() error { + err := rows.Rows.Err() + rows.unit.recordError(err) + return err +} + +func (rows unitRows) Scan(dest ...any) error { + err := rows.Rows.Scan(dest...) + rows.unit.recordError(err) + return err +} + +func (rows unitRows) Close() { + rows.Rows.Close() + rows.unit.recordError(rows.Rows.Err()) +} + +func (db unitDB) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { + if err := db.unit.active(); err != nil { + return failedBatch{err: err} + } + return unitBatch{BatchResults: db.unit.tx.SendBatch(ctx, batch), unit: db.unit} +} + +type failedBatch struct{ err error } + +func (b failedBatch) Exec() (pgconn.CommandTag, error) { return pgconn.CommandTag{}, b.err } +func (b failedBatch) Query() (pgx.Rows, error) { return nil, b.err } +func (b failedBatch) QueryRow() pgx.Row { return errorRow(b) } +func (b failedBatch) Close() error { return b.err } + +type unitBatch struct { + pgx.BatchResults + unit *Unit +} + +func (b unitBatch) Exec() (pgconn.CommandTag, error) { + tag, err := b.BatchResults.Exec() + b.unit.recordError(err) + return tag, err +} + +func (b unitBatch) Query() (pgx.Rows, error) { + rows, err := b.BatchResults.Query() + b.unit.recordError(err) + if err != nil { + return nil, err + } + return unitRows{Rows: rows, unit: b.unit}, nil +} + +func (b unitBatch) QueryRow() pgx.Row { + return unitRow{row: b.BatchResults.QueryRow(), unit: b.unit} +} + +func (b unitBatch) Close() error { + err := b.BatchResults.Close() + b.unit.recordError(err) + return err +} diff --git a/internal/storage/internal/agentexecution/evidence_test.go b/internal/storage/internal/agentexecution/evidence_test.go new file mode 100644 index 000000000..f6ef0e45c --- /dev/null +++ b/internal/storage/internal/agentexecution/evidence_test.go @@ -0,0 +1,171 @@ +package agentexecution + +import ( + "encoding/json" + "math" + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/toolcatalog" + "github.com/stretchr/testify/require" +) + +func TestBindToolCallsUsesProviderEnvelopeAsCanonicalSource(t *testing.T) { + id := uuid.New() + envelope := modelenvelope.ResponseEnvelope{ + Normalized: modelenvelope.ResponseNormalized{ + Content: []modelenvelope.ResponsePart{ + {Type: modelenvelope.ResponsePartTypeText, Text: "before"}, + { + Type: modelenvelope.ResponsePartTypeToolCall, + ProviderCallID: "call_1", + ToolName: "lookup_customer", + ToolInput: json.RawMessage(`{"email":"ada@example.com"}`), + }, + }, + }, + } + _, calls, err := outputParts(envelope, []ToolProposal{{ + ID: id, + ProviderCallID: "call_1", + Type: toolcatalog.ToolTypeCustom, + }}) + if err != nil { + t.Fatalf("bind tool calls: %v", err) + } + if len(calls) != 1 || + calls[0].ID != id || + calls[0].ProviderCallID != "call_1" || + calls[0].Name != "lookup_customer" || + string(calls[0].Input) != `{"email":"ada@example.com"}` || + calls[0].Type != toolcatalog.ToolTypeCustom { + t.Fatalf("bound tool call = %+v", calls) + } +} + +func TestBindToolCallsRequiresExactProviderCallCoverage(t *testing.T) { + envelope := modelenvelope.ResponseEnvelope{ + Normalized: modelenvelope.ResponseNormalized{ + Content: []modelenvelope.ResponsePart{{ + Type: modelenvelope.ResponsePartTypeToolCall, + ProviderCallID: "call_1", + ToolName: "run_command", + ToolInput: json.RawMessage(`{"command":"true"}`), + }}, + }, + } + for _, bindings := range [][]ToolProposal{ + {{ + ProviderCallID: "call_other", + Type: toolcatalog.ToolTypeBuiltIn, + }}, + { + { + ProviderCallID: "call_1", + Type: toolcatalog.ToolTypeBuiltIn, + }, + { + ProviderCallID: "call_extra", + Type: toolcatalog.ToolTypeBuiltIn, + }, + }, + } { + if _, _, err := outputParts(envelope, bindings); err == nil { + t.Fatalf("bindings %+v did not fail", bindings) + } + } +} + +func TestBindToolCallsRejectsInvalidToolInput(t *testing.T) { + for _, input := range []json.RawMessage{ + json.RawMessage(`{"command":`), + json.RawMessage(`null`), + json.RawMessage(`[]`), + json.RawMessage(`"command"`), + } { + envelope := modelenvelope.ResponseEnvelope{ + Normalized: modelenvelope.ResponseNormalized{ + Content: []modelenvelope.ResponsePart{{ + Type: modelenvelope.ResponsePartTypeToolCall, + ProviderCallID: "call_1", + ToolName: "run_command", + ToolInput: input, + }}, + }, + } + _, _, err := outputParts(envelope, []ToolProposal{{ + ID: uuid.New(), + ProviderCallID: "call_1", + Type: toolcatalog.ToolTypeBuiltIn, + }}) + if err == nil { + t.Fatalf("outputParts accepted input %s", input) + } + } +} + +func TestModelUsageForStorageNormalizesProviderUsage(t *testing.T) { + usage := normalizedUsage(modelenvelope.Usage{ + InputTokens: 20, + OutputTokens: 7, + ReasoningTokens: 3, + CacheReadTokens: 8, + CacheWriteTokens: 2, + }) + if usage != (modelenvelope.Usage{ + InputTokens: 20, + UncachedInputTokens: 10, + OutputTokens: 7, + ReasoningTokens: 3, + CacheReadTokens: 8, + CacheWriteTokens: 2, + }) { + t.Fatalf("normalized usage = %+v", usage) + } +} + +func TestModelUsageForStorageDropsInvalidOrUnrepresentableUsage(t *testing.T) { + tests := []modelenvelope.Usage{ + {InputTokens: -1}, + {InputTokens: math.MaxInt32 + 1}, + {InputTokens: 10, OutputTokens: 2, ReasoningTokens: 3}, + {InputTokens: 10, CacheReadTokens: 11}, + } + for _, usage := range tests { + if got := normalizedUsage(usage); got != (modelenvelope.Usage{}) { + t.Fatalf("invalid usage %+v normalized to %+v", usage, got) + } + } +} +func TestEvidenceIdentityAndCost(t *testing.T) { + for _, e := range []ModelEvidence{{RequestID: "req"}, + {ResponseID: "response"}, + {Cost: "0.01"}, + {Usage: modelenvelope.Usage{InputTokens: 1}}, + {APIFormat: "openai-responses"}, + {APIFormat: "openai-responses", + APIVariant: "default", + Cost: "invalid"}} { + _, err := evidenceParams(e) + require.Error(t, err) + } + for _, cost := range []modelenvelope.ProviderReportedCostUSD{"", "0", "0.0000125", "100"} { + columns, err := evidenceParams( + ModelEvidence{ + APIFormat: "openai-responses", + APIVariant: "default", + Cost: cost, + Usage: modelenvelope.Usage{InputTokens: 10, UncachedInputTokens: 10, OutputTokens: 4}, + }, + ) + require.NoError(t, err) + require.EqualValues(t, cost, valueOrZero(columns.Cost)) + require.EqualValues(t, 10, *columns.InputTokens) + require.EqualValues(t, 10, *columns.UncachedTokens) + require.EqualValues(t, 4, *columns.OutputTokens) + require.Nil(t, columns.CacheReadTokens) + require.Nil(t, columns.CacheWriteTokens) + require.Nil(t, columns.ReasoningTokens) + } +} diff --git a/internal/storage/internal/agentexecution/fixture_integration_test.go b/internal/storage/internal/agentexecution/fixture_integration_test.go new file mode 100644 index 000000000..9c2596440 --- /dev/null +++ b/internal/storage/internal/agentexecution/fixture_integration_test.go @@ -0,0 +1,82 @@ +//go:build integration + +package agentexecution_test + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/omnara-ai/omnara/internal/storage" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + "github.com/omnara-ai/omnara/internal/testutil/storagefixture" + "github.com/stretchr/testify/require" +) + +func migratedExecutionFixture(t *testing.T) executionFixture { + t.Helper() + return seedExecutionFixture(t, integrationdb.OpenMigratedPool(t, t.Context(), "../../../../migrations")) +} + +func seedExecutionFixture(t *testing.T, pool *pgxpool.Pool) executionFixture { + t.Helper() + ctx := t.Context() + ids := storagefixture.ProjectIDs{ + OrgID: uuid.New(), + ProjectID: uuid.New(), + ProviderAdminUserID: uuid.New(), + ProviderSecretID: uuid.New(), + ProviderSecretVersionID: uuid.New(), + ProviderConfigID: uuid.New(), + } + now := time.Now().UTC().Truncate(time.Microsecond) + storagefixture.SeedProject(t, ctx, pool, ids, now) + store := storage.NewStore(pool) + config := storagefixture.SeedAgentConfig( + t, + ctx, + store.Models(), + store.Execution(), + ids.OrgID, + ids.ProjectID, + "instruction: test\nmodel:\n provider_config: openai-prod\n name: test\n", + ) + var revision uuid.UUID + require.NoError( + t, + pool.QueryRow(ctx, + `SELECT current_revision_id FROM configured_models WHERE id=$1`, + config.ConfiguredModelID). + Scan(&revision), + ) + return executionFixture{ + pool: pool, + org: ids.OrgID, + project: ids.ProjectID, + config: config.ID, + revision: revision, + now: now, + } +} + +func (f executionFixture) create(t *testing.T, id, parent uuid.UUID) uuid.UUID { + t.Helper() + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = u.Rollback(t.Context()) }() + key := "" + if parent != uuid.Nil { + key = id.String() + require.NoError(t, u.LockAgentRefs(t.Context(), + []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: parent}}, agentexecution.LifecycleAuthority{})) + } + _, err = u.CreateAgent(t.Context(), agentexecution.CreateAgentInput{ID: id, ProjectID: f.project, + ConfigID: f.config, ParentID: parent, SubagentKey: key, Name: "Agent"}) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "create fixture")) + return id +} diff --git a/internal/storage/internal/agentexecution/handle.go b/internal/storage/internal/agentexecution/handle.go new file mode 100644 index 000000000..e88490c2f --- /dev/null +++ b/internal/storage/internal/agentexecution/handle.go @@ -0,0 +1,47 @@ +package agentexecution + +import ( + "fmt" + + "github.com/google/uuid" +) + +type Handle struct { + unit *Unit + route AgentRoute + valid bool + created bool + mutation *executionMutation +} + +func (u *Unit) Agent(route AgentRoute) (*Handle, error) { + if err := u.active(); err != nil { + return nil, err + } + h, found := u.handles[route.AgentID] + if !found || h.route != route { + return nil, fmt.Errorf("%w: agent is not locked in this unit", ErrLockPlan) + } + return h, nil +} + +func (h *Handle) Route() (AgentRoute, error) { + if !h.valid || h.unit.handles[h.route.AgentID] != h { + return AgentRoute{}, ErrUnitClosed + } + if err := h.unit.active(); err != nil { + return AgentRoute{}, err + } + return h.route, nil +} + +func (u *Unit) Handle(projectID, agentID uuid.UUID) (*Handle, error) { + if err := u.active(); err != nil { + return nil, err + } + h := u.handles[agentID] + if h == nil || h.route.ProjectID != projectID { + return nil, fmt.Errorf("%w: agent is not locked in this unit", ErrLockPlan) + } + return h, nil +} diff --git a/internal/storage/internal/agentexecution/history_test.go b/internal/storage/internal/agentexecution/history_test.go new file mode 100644 index 000000000..01b7f3046 --- /dev/null +++ b/internal/storage/internal/agentexecution/history_test.go @@ -0,0 +1,306 @@ +package agentexecution_test + +import ( + "fmt" + "math/rand/v2" + "testing" + "time" + + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +type historyOutput struct { + output agentexecution.ModelOutput + tools []historyTool +} + +type historyTool struct { + state string + kind string + resultSequence int64 + completedAt time.Time +} + +type executionHistory struct { + view agentexecution.ExecutionView + contexts []agentexecution.ModelContext + outputs []historyOutput +} + +func (h executionHistory) reconstruct() (agentexecution.ExecutionView, error) { + view := h.view + view.NormalContext, view.CompactionContext, view.ToolBatch = nil, nil, nil + view.MaxNormalInputSequence, view.MaxContextInputSequence = 0, 0 + started := 0 + identities := make(map[string]bool) + for _, context := range h.contexts { + key := fmt.Sprintf("%s/%d/%d/%d", context.Operation, context.InputEventSequence, + context.SourceEventSequenceEnd, context.Attempt) + if identities[key] { + return view, fmt.Errorf("%w: duplicate context identity", agentexecution.ErrInvalidState) + } + identities[key] = true + if context.State == agentexecution.ContextStarted { + started++ + } + view.MaxContextInputSequence = max(view.MaxContextInputSequence, context.InputEventSequence) + if context.Operation == agentexecution.OperationNormal { + view.MaxNormalInputSequence = max(view.MaxNormalInputSequence, context.InputEventSequence) + if context.TurnID == view.Turn.ID && (view.NormalContext == nil || + context.InputEventSequence > view.NormalContext.InputEventSequence || + (context.InputEventSequence == view.NormalContext.InputEventSequence && + context.Attempt > view.NormalContext.Attempt)) { + view.NormalContext = &context + } + } + } + if started > 1 { + return view, fmt.Errorf("%w: concurrent started contexts", agentexecution.ErrInvalidState) + } + for _, context := range h.contexts { + if context.Operation != agentexecution.OperationCompaction || context.TurnID != view.Turn.ID { + continue + } + if view.NormalContext == nil || context.InputEventSequence > view.NormalContext.InputEventSequence { + return view, fmt.Errorf("%w: orphan compaction", agentexecution.ErrInvalidState) + } + if context.InputEventSequence != view.NormalContext.InputEventSequence { + continue + } + if view.CompactionContext == nil || + context.SourceEventSequenceEnd < view.CompactionContext.SourceEventSequenceEnd || + (context.SourceEventSequenceEnd == view.CompactionContext.SourceEventSequenceEnd && + context.Attempt > view.CompactionContext.Attempt) { + view.CompactionContext = &context + } + } + for _, output := range h.outputs { + if output.output.Context.TurnID != view.Turn.ID || output.output.Event.Sequence < view.StopSequence || + len(output.tools) == 0 { + continue + } + batch := agentexecution.ToolBatch{Output: output.output} + completion := agentexecution.BatchCompletion{} + for _, tool := range output.tools { + if tool.state != "completed" { + batch.HasIncomplete = true + } + if tool.state == "awaiting_authorization" || + (tool.state == "ready" && (tool.kind == "built_in" || tool.kind == "mcp")) { + batch.HasRunnable = true + } + if tool.state == "completed" { + if tool.resultSequence <= output.output.Event.Sequence || tool.completedAt.IsZero() { + return view, fmt.Errorf("%w: missing result authority", agentexecution.ErrInvalidState) + } + completion.LastResultSequence = max(completion.LastResultSequence, tool.resultSequence) + if tool.completedAt.After(completion.ReadyAt) { + completion.ReadyAt = tool.completedAt + } + } + } + if !batch.HasIncomplete { + batch.Completion = &completion + consumed := false + for _, later := range h.outputs { + if later.output.Context.TurnID == output.output.Context.TurnID && + later.output.Event.Sequence > completion.LastResultSequence && + later.output.Context.InputEventSequence >= completion.LastResultSequence { + consumed = true + } + } + if consumed { + continue + } + } + if view.ToolBatch != nil { + return view, fmt.Errorf("%w: multiple unconsumed batches", agentexecution.ErrInvalidState) + } + view.ToolBatch = &batch + } + return view, nil +} + +func TestSingleBatchHistoryTrace(t *testing.T) { + t.Parallel() + for _, seed := range []uint64{1, 7, 42, 101} { + t.Run(fmt.Sprint(seed), func(t *testing.T) { + t.Parallel() + rng := rand.New(rand.NewPCG(seed, seed+1)) + base := baseView() + history := executionHistory{view: base, contexts: []agentexecution.ModelContext{*base.NormalContext}} + previousResult := int64(0) + for cycle := range 16 { + context := history.contexts[len(history.contexts)-1] + out := agentexecution.ModelOutput{ID: identity(byte(30 + cycle)), EventID: identity(byte(60 + cycle)), + Context: context, Event: agentexecution.EventBoundary{ + Sequence: context.InputEventSequence + 1, Time: decisionTime().Add(-time.Second), + }} + history.outputs = append(history.outputs, historyOutput{output: out, + tools: []historyTool{{state: "ready", kind: "built_in"}, {state: "ready", kind: "mcp"}}}) + view, err := history.reconstruct() + require.NoError(t, err) + require.Equal(t, out.ID, view.ToolBatch.Output.ID) + selection, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, agentexecution.WorkTool, selection.Work) + for step, index := range rng.Perm(2) { + tool := &history.outputs[cycle].tools[index] + tool.state, tool.resultSequence = "completed", out.Event.Sequence+int64(step)+1 + tool.completedAt = decisionTime().Add(-time.Duration(2-step) * time.Millisecond) + view, err = history.reconstruct() + require.NoError(t, err) + selection, err = agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + if step == 0 { + require.Equal(t, agentexecution.WorkTool, selection.Work) + } else { + require.Equal(t, agentexecution.OriginTools, selection.Model.Origin) + } + } + previousResult = view.ToolBatch.Completion.LastResultSequence + next := context + next.Attempt = 1 + next.ID, next.InputEventSequence, next.State = identity(byte(90+cycle)), previousResult, + agentexecution.ContextStarted + history.contexts = append(history.contexts, next) + view, err = history.reconstruct() + require.NoError(t, err) + require.Equal(t, out.ID, view.ToolBatch.Output.ID) + deadline := decisionTime().Add(time.Hour) + next.State, next.Recovery, next.RetryAt = agentexecution.ContextFailed, agentexecution.RecoveryRetry, &deadline + history.contexts[len(history.contexts)-1] = next + view, err = history.reconstruct() + require.NoError(t, err) + selection, err = agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, agentexecution.OriginRetry, selection.Model.Origin) + require.Equal(t, agentexecution.WaitModelDeadline, selection.Wait) + require.Equal(t, out.ID, view.ToolBatch.Output.ID) + next.ID, next.Attempt, next.State = identity(byte(120+cycle)), 2, agentexecution.ContextSucceeded + next.Recovery, next.RetryAt = agentexecution.RecoveryNone, nil + history.contexts = append(history.contexts, next) + } + last := history.contexts[len(history.contexts)-1] + history.outputs = append(history.outputs, historyOutput{output: agentexecution.ModelOutput{ + ID: identity(240), EventID: identity(241), Context: last, + Event: agentexecution.EventBoundary{Sequence: previousResult + 1, Time: decisionTime()}, + }}) + view, err := history.reconstruct() + require.NoError(t, err) + require.Nil(t, view.ToolBatch) + }) + } +} + +func TestGoverningContextHistoryTrace(t *testing.T) { + t.Parallel() + view := compactionView() + normal, compaction := *view.NormalContext, *view.CompactionContext + compaction.State, compaction.Recovery = agentexecution.ContextFailed, agentexecution.RecoveryRetry + deadline := decisionTime() + compaction.RetryAt = &deadline + history := executionHistory{view: view, contexts: []agentexecution.ModelContext{normal, compaction}} + for attempt := int32(2); attempt <= 9; attempt++ { + next := compaction + next.ID, next.Attempt = identity(byte(10+attempt)), attempt + history.contexts = append(history.contexts, next) + } + head, err := history.reconstruct() + require.NoError(t, err) + require.EqualValues(t, 9, head.CompactionContext.Attempt) + selection, err := agentexecution.SelectNext(head, decisionTime()) + require.NoError(t, err) + require.Equal(t, head.CompactionContext.ID, selection.Model.SourceContextID) + for end := int64(8); end > 0; end-- { + next := compaction + next.ID, next.Attempt, next.SourceEventSequenceEnd = identity(byte(30+end)), 1, end + history.contexts = append(history.contexts, next) + head, err = history.reconstruct() + require.NoError(t, err) + require.Equal(t, next.ID, head.CompactionContext.ID) + } + terminal := &history.contexts[len(history.contexts)-1] + terminal.State, terminal.Recovery, terminal.RetryAt = agentexecution.ContextSucceeded, agentexecution.RecoveryNone, nil + head, err = history.reconstruct() + require.NoError(t, err) + selection, err = agentexecution.SelectNext(head, decisionTime()) + require.NoError(t, err) + require.Nil(t, selection.Model) + require.False(t, selection.TurnContinuable) + next := normal + next.ID, next.InputEventSequence, next.State = identity(80), 15, agentexecution.ContextStarted + next.Recovery = agentexecution.RecoveryNone + history.contexts = append(history.contexts, next) + head, err = history.reconstruct() + require.NoError(t, err) + require.Equal(t, next.ID, head.NormalContext.ID) + require.Nil(t, head.CompactionContext) +} + +func TestHistoryRejectsMalformedObligations(t *testing.T) { + t.Parallel() + for _, malformed := range []string{"duplicate context", "two live contexts", "orphan compaction", + "two incomplete batches", "two unconsumed complete batches", "missing result"} { + t.Run(malformed, func(t *testing.T) { + t.Parallel() + view := baseView() + history := executionHistory{view: view, contexts: []agentexecution.ModelContext{*view.NormalContext}} + switch malformed { + case "duplicate context": + other := *view.NormalContext + other.ID = identity(90) + history.contexts = append(history.contexts, other) + case "two live contexts": + history.contexts[0].State = agentexecution.ContextStarted + other := history.contexts[0] + other.ID, other.InputEventSequence = identity(90), 11 + history.contexts = append(history.contexts, other) + case "orphan compaction": + other := *compactionView().CompactionContext + other.InputEventSequence = 11 + history.contexts = append(history.contexts, other) + case "two incomplete batches", "two unconsumed complete batches": + for i := range 2 { + o := historyOutput{output: output(view), tools: []historyTool{{state: "waiting", kind: "custom"}}} + o.output.ID, o.output.EventID, o.output.Event.Sequence = identity(byte(90+i)), identity(byte(92+i)), 20+int64(i) + if malformed == "two unconsumed complete batches" { + o.tools[0] = historyTool{state: "completed", resultSequence: 30 + int64(i), completedAt: decisionTime()} + } + history.outputs = append(history.outputs, o) + } + case "missing result": + history.outputs = []historyOutput{{output: output(view), tools: []historyTool{{state: "completed"}}}} + } + _, err := history.reconstruct() + require.ErrorIs(t, err, agentexecution.ErrInvalidState) + }) + } +} + +func TestToolCustodyFacts(t *testing.T) { + t.Parallel() + for _, kind := range []string{"built_in", "mcp", "custom"} { + for _, state := range []string{"awaiting_authorization", "ready", "running", "waiting", "completed"} { + t.Run(kind+"/"+state, func(t *testing.T) { + t.Parallel() + view := baseView() + history := executionHistory{view: view, contexts: []agentexecution.ModelContext{*view.NormalContext}, + outputs: []historyOutput{{output: output(view), tools: []historyTool{{ + state: state, kind: kind, resultSequence: 21, completedAt: decisionTime(), + }}}}} + view, err := history.reconstruct() + require.NoError(t, err) + got, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + runnable := state == "awaiting_authorization" || (state == "ready" && kind != "custom") + require.Equal(t, runnable, got.Work == agentexecution.WorkTool) + require.Equal(t, state != "completed", got.IncompleteTools) + if !runnable && state != "completed" { + require.Equal(t, agentexecution.WaitToolBatch, got.Wait) + } + }) + } + } +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/batch.go b/internal/storage/internal/agentexecution/internal/executiondb/batch.go new file mode 100644 index 000000000..4614c980c --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/batch.go @@ -0,0 +1,548 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: batch.go + +package executiondb + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" +) + +var ( + ErrBatchAlreadyClosed = errors.New("batch already closed") +) + +const loadExecutionFacts = `-- name: LoadExecutionFacts :batchmany +WITH scope AS MATERIALIZED ( + SELECT a.id,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now,t.id AS turn_id, + coalesce(e.sequence,0)::bigint AS semantic_sequence,coalesce(e.created_at,'epoch'::timestamptz)::timestamptz AS semantic_time, + EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='steering' AND i.state='received') AS steering, + EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='queued' AND i.state='received') AS queued, + coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event + ORDER BY x.sequence LIMIT 1),0)::bigint AS first_opening_sequence, + coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event + ORDER BY x.sequence DESC LIMIT 1),0)::bigint AS last_opening_sequence, + coalesce((SELECT x.sequence FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id + WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event AND i.input_kind='content' + ORDER BY x.sequence LIMIT 1),0)::bigint AS first_content_sequence + FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=$3 + LEFT JOIN agent_events e ON e.agent_id=a.id AND e.id=t.latest_semantic_event_id + WHERE a.id=$1 AND a.project_id=$4 +), outputs AS MATERIALIZED ( + SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at + FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id + WHERE o.agent_id=$1 AND o.id IN ($5::uuid,$6::uuid) +), batch AS MATERIALIZED ( + SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 + AND c.model_output_id=$5) AS has_tools, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 + AND c.model_output_id=$5 AND c.state<>'completed') AS incomplete, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 + AND c.model_output_id=$5 AND + (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable + WHERE $5::uuid IS NOT NULL +), boundary AS MATERIALIZED ( + SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$1 + AND e.event_kind='model_output' AND e.sequence<=(SELECT last_opening_sequence FROM scope) + ORDER BY e.sequence DESC LIMIT 1),0),$7::bigint)::bigint AS sequence + FROM scope WHERE scope.state='active' AND $8::bigintb.sequence AND e.sequence<=(SELECT last_opening_sequence FROM scope) + JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id +), opening AS ( + SELECT id,sequence,created_at FROM unanswered + UNION ALL + SELECT i.id,e.sequence,e.created_at FROM boundary b + JOIN agent_events e ON e.agent_id=$1 AND e.turn_id=$3 AND e.is_opening_event + AND e.sequence<=(SELECT last_opening_sequence FROM scope) + JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id + WHERE NOT EXISTS(SELECT 1 FROM unanswered) +), facts AS ( +SELECT 'context'::text AS kind, + c.id AS id, + c.turn_id AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + 0::bigint AS sequence, + 'epoch'::timestamptz AS event_time, + c.operation_kind AS operation, + c.attempt_number AS attempt, + c.input_event_sequence AS watermark, + coalesce(c.source_event_sequence_end,0)::bigint AS source_end, + c.state AS state, + coalesce(c.recovery_kind,'')::text AS recovery, + c.retry_at AS retry_at, + c.opening_input_ids AS opening_ids, + c.opening_event_sequence AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM model_call_contexts c WHERE c.agent_id=$1 AND c.id IN ( + SELECT $9::uuid UNION SELECT $10::uuid + UNION SELECT model_call_context_id FROM outputs) +UNION ALL +SELECT 'output'::text AS kind, + o.id AS id, + '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, + o.model_call_context_id AS context_id, + o.event_id AS event_id, + o.sequence AS sequence, + o.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + '{}'::uuid[] AS opening_ids, + 0::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + o.stop_reason AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM outputs o +UNION ALL +SELECT 'batch'::text AS kind, + '00000000-0000-0000-0000-000000000000'::uuid AS id, + '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + coalesce(r.last_sequence,0)::bigint AS sequence, + coalesce(r.ready_at,'epoch'::timestamptz)::timestamptz AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + '{}'::uuid[] AS opening_ids, + 0::bigint AS opening_sequence, + b.has_tools AS has_tools, + b.incomplete AS incomplete, + b.runnable AS runnable, + coalesce(r.calls,0)::bigint AS calls, + coalesce(r.results,0)::bigint AS results, + coalesce(r.events,0)::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM batch b LEFT JOIN LATERAL ( + SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, + max(e.sequence)::bigint AS last_sequence,max(r.completed_at)::timestamptz AS ready_at + FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id + LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id + WHERE c.agent_id=$1 AND c.model_output_id=$5 AND NOT b.incomplete +) r ON NOT b.incomplete +UNION ALL +SELECT 'config'::text AS kind, + i.id AS id, + e.turn_id AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + e.sequence AS sequence, + e.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + i.opening_input_ids AS opening_ids, + coalesce(i.opening_event_sequence,0)::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id +WHERE i.agent_id=$1 AND i.id=$11 AND i.input_kind='config_change' AND i.state='resolved' +UNION ALL +SELECT 'checkpoint'::text AS kind, + c.id AS id, + e.turn_id AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + e.sequence AS sequence, + c.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + c.opening_input_ids AS opening_ids, + coalesce(c.opening_event_sequence,0)::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=$1 AND c.id=$12 +UNION ALL +SELECT 'opening'::text AS kind, + o.id AS id, + '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + o.sequence AS sequence, + o.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + '{}'::uuid[] AS opening_ids, + 0::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM opening o +UNION ALL +SELECT 'scope'::text AS kind,s.id,coalesce(s.turn_id,'00000000-0000-0000-0000-000000000000'::uuid), + '00000000-0000-0000-0000-000000000000'::uuid,'00000000-0000-0000-0000-000000000000'::uuid, + s.semantic_sequence,s.semantic_time,''::text,0::integer,0::bigint,0::bigint,s.state,''::text,NULL::timestamptz, + '{}'::uuid[],0::bigint,false,false,false,0::bigint,0::bigint,0::bigint,''::text, + s.root_agent_id,s.database_now,s.steering,s.queued,s.first_opening_sequence,s.last_opening_sequence,s.first_content_sequence +FROM scope s + ) +SELECT f.kind,f.id,f.turn_id,f.context_id,f.event_id,f.sequence,f.event_time,f.operation,f.attempt, + f.watermark,f.source_end,f.state,f.recovery,f.retry_at,f.opening_ids,f.opening_sequence, + f.has_tools,f.incomplete,f.runnable,f.calls,f.results,f.events,f.stop_reason,f.root_agent_id, + f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence, + coalesce(w.org_id,'00000000-0000-0000-0000-000000000000'::uuid)::uuid AS webhook_org_id, + coalesce(w.events,'[]'::jsonb)::jsonb AS webhook_events +FROM facts f LEFT JOIN LATERAL ( + SELECT p.org_id,c.compiled_definition->'event_webhook'->'events' AS events + FROM agents a JOIN agent_configs c ON c.project_id=a.project_id AND c.id=a.current_config_id + JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id + WHERE a.id=$1 AND p.deleted_at IS NULL AND o.deleted_at IS NULL + AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' + AND $2::boolean AND f.kind='scope' +) w ON true ORDER BY f.kind,f.sequence +` + +type LoadExecutionFactsBatchResults struct { + br pgx.BatchResults + tot int + closed bool +} + +type LoadExecutionFactsParams struct { + AgentID uuid.UUID + CaptureWebhook bool + TurnID *uuid.UUID + ProjectID uuid.UUID + ToolOutputID *uuid.UUID + LimitOutputID *uuid.UUID + StopSequence int64 + MaxContextInputSequence int64 + NormalContextID *uuid.UUID + CompactionContextID *uuid.UUID + ConfigID *uuid.UUID + CheckpointID *uuid.UUID +} + +type LoadExecutionFactsRow struct { + Kind string + ID uuid.UUID + TurnID uuid.UUID + ContextID uuid.UUID + EventID uuid.UUID + Sequence int64 + EventTime time.Time + Operation string + Attempt int32 + Watermark int64 + SourceEnd int64 + State string + Recovery string + RetryAt *time.Time + OpeningIds []uuid.UUID + OpeningSequence int64 + HasTools bool + Incomplete bool + Runnable bool + Calls int64 + Results int64 + Events int64 + StopReason string + RootAgentID uuid.UUID + DatabaseNow time.Time + Steering bool + Queued bool + FirstOpeningSequence int64 + LastOpeningSequence int64 + FirstContentSequence int64 + WebhookOrgID uuid.UUID + WebhookEvents json.RawMessage +} + +func (q *Queries) LoadExecutionFacts(ctx context.Context, db DBTX, arg []LoadExecutionFactsParams) *LoadExecutionFactsBatchResults { + batch := &pgx.Batch{} + for _, a := range arg { + vals := []interface{}{ + a.AgentID, + a.CaptureWebhook, + a.TurnID, + a.ProjectID, + a.ToolOutputID, + a.LimitOutputID, + a.StopSequence, + a.MaxContextInputSequence, + a.NormalContextID, + a.CompactionContextID, + a.ConfigID, + a.CheckpointID, + } + batch.Queue(loadExecutionFacts, vals...) + } + br := db.SendBatch(ctx, batch) + return &LoadExecutionFactsBatchResults{br, len(arg), false} +} + +func (b *LoadExecutionFactsBatchResults) Query(f func(int, []LoadExecutionFactsRow, error)) { + defer b.br.Close() + for t := 0; t < b.tot; t++ { + items := []LoadExecutionFactsRow{} + if b.closed { + if f != nil { + f(t, items, ErrBatchAlreadyClosed) + } + continue + } + err := func() error { + rows, err := b.br.Query() + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var i LoadExecutionFactsRow + if err := rows.Scan( + &i.Kind, + &i.ID, + &i.TurnID, + &i.ContextID, + &i.EventID, + &i.Sequence, + &i.EventTime, + &i.Operation, + &i.Attempt, + &i.Watermark, + &i.SourceEnd, + &i.State, + &i.Recovery, + &i.RetryAt, + &i.OpeningIds, + &i.OpeningSequence, + &i.HasTools, + &i.Incomplete, + &i.Runnable, + &i.Calls, + &i.Results, + &i.Events, + &i.StopReason, + &i.RootAgentID, + &i.DatabaseNow, + &i.Steering, + &i.Queued, + &i.FirstOpeningSequence, + &i.LastOpeningSequence, + &i.FirstContentSequence, + &i.WebhookOrgID, + &i.WebhookEvents, + ); err != nil { + return err + } + items = append(items, i) + } + return rows.Err() + }() + if f != nil { + f(t, items, err) + } + } +} + +func (b *LoadExecutionFactsBatchResults) Close() error { + b.closed = true + return b.br.Close() +} + +const publishExecution = `-- name: PublishExecution :batchone +WITH scope AS MATERIALIZED ( + SELECT a.id AS agent_id,$1::timestamptz AS logical_ready_at + FROM agents a WHERE a.id=$2 AND ($3::boolean OR + EXISTS(SELECT 1 FROM agent_execution_state h WHERE h.agent_id=a.id)) +), updated AS ( +INSERT INTO agent_execution_state AS h(agent_id,current_turn_id,stop_sequence,answered_through_sequence,max_normal_input_sequence,max_context_input_sequence,normal_context_id,compaction_context_id,pending_tool_output_id,pending_output_limit_id,pending_config_input_id,pending_checkpoint_id,turn_continuable,incomplete_tools,logical_ready_at) +SELECT scope.agent_id,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,scope.logical_ready_at FROM scope +ON CONFLICT(agent_id) DO UPDATE SET +current_turn_id=excluded.current_turn_id, +stop_sequence=excluded.stop_sequence, +answered_through_sequence=excluded.answered_through_sequence, +max_normal_input_sequence=excluded.max_normal_input_sequence, +max_context_input_sequence=excluded.max_context_input_sequence, +normal_context_id=excluded.normal_context_id, +compaction_context_id=excluded.compaction_context_id, +pending_tool_output_id=excluded.pending_tool_output_id, +pending_output_limit_id=excluded.pending_output_limit_id, +pending_config_input_id=excluded.pending_config_input_id, +pending_checkpoint_id=excluded.pending_checkpoint_id, +turn_continuable=excluded.turn_continuable, +incomplete_tools=excluded.incomplete_tools, +logical_ready_at=excluded.logical_ready_at +WHERE (h.current_turn_id,h.stop_sequence,h.answered_through_sequence,h.max_normal_input_sequence,h.max_context_input_sequence,h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id,h.pending_config_input_id,h.pending_checkpoint_id,h.turn_continuable,h.incomplete_tools,h.logical_ready_at) IS DISTINCT FROM +(excluded.current_turn_id,excluded.stop_sequence,excluded.answered_through_sequence,excluded.max_normal_input_sequence,excluded.max_context_input_sequence,excluded.normal_context_id,excluded.compaction_context_id,excluded.pending_tool_output_id,excluded.pending_output_limit_id,excluded.pending_config_input_id,excluded.pending_checkpoint_id,excluded.turn_continuable,excluded.incomplete_tools,excluded.logical_ready_at) +RETURNING h.agent_id), desired AS MATERIALIZED ( + SELECT a.id,h.logical_ready_at AS ready_at FROM scope h JOIN agents a ON a.id=h.agent_id + WHERE a.id=$2 AND a.state='active' AND h.logical_ready_at IS NOT NULL + AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id) +), removed AS ( + DELETE FROM agent_wakeups WHERE agent_id=$2 AND NOT EXISTS(SELECT 1 FROM desired) +), published AS ( +INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) +SELECT id,ready_at,statement_timestamp() FROM desired +ON CONFLICT(agent_id) DO UPDATE SET ready_at=CASE + WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() + THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END, + updated_at=statement_timestamp() +WHERE agent_wakeups.ready_at IS DISTINCT FROM CASE + WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() + THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END +RETURNING agent_id +) SELECT agent_id FROM scope +` + +type PublishExecutionBatchResults struct { + br pgx.BatchResults + tot int + closed bool +} + +type PublishExecutionParams struct { + LogicalReadyAt *time.Time + AgentID uuid.UUID + InsertHead bool + CurrentTurnID *uuid.UUID + StopSequence int64 + AnsweredThroughSequence int64 + MaxNormalInputSequence int64 + MaxContextInputSequence int64 + NormalContextID *uuid.UUID + CompactionContextID *uuid.UUID + PendingToolOutputID *uuid.UUID + PendingOutputLimitID *uuid.UUID + PendingConfigInputID *uuid.UUID + PendingCheckpointID *uuid.UUID + TurnContinuable bool + IncompleteTools bool +} + +func (q *Queries) PublishExecution(ctx context.Context, db DBTX, arg []PublishExecutionParams) *PublishExecutionBatchResults { + batch := &pgx.Batch{} + for _, a := range arg { + vals := []interface{}{ + a.LogicalReadyAt, + a.AgentID, + a.InsertHead, + a.CurrentTurnID, + a.StopSequence, + a.AnsweredThroughSequence, + a.MaxNormalInputSequence, + a.MaxContextInputSequence, + a.NormalContextID, + a.CompactionContextID, + a.PendingToolOutputID, + a.PendingOutputLimitID, + a.PendingConfigInputID, + a.PendingCheckpointID, + a.TurnContinuable, + a.IncompleteTools, + } + batch.Queue(publishExecution, vals...) + } + br := db.SendBatch(ctx, batch) + return &PublishExecutionBatchResults{br, len(arg), false} +} + +func (b *PublishExecutionBatchResults) QueryRow(f func(int, uuid.UUID, error)) { + defer b.br.Close() + for t := 0; t < b.tot; t++ { + var agent_id uuid.UUID + if b.closed { + if f != nil { + f(t, agent_id, ErrBatchAlreadyClosed) + } + continue + } + row := b.br.QueryRow() + err := row.Scan(&agent_id) + if f != nil { + f(t, agent_id, err) + } + } +} + +func (b *PublishExecutionBatchResults) Close() error { + b.closed = true + return b.br.Close() +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_agents.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_agents.sql.go new file mode 100644 index 000000000..00a41dcc5 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_agents.sql.go @@ -0,0 +1,181 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_agents.sql + +package executiondb + +import ( + "context" + + "github.com/google/uuid" +) + +const archiveExecutionAgent = `-- name: ArchiveExecutionAgent :execrows +UPDATE agents SET state='archived',archived_at=coalesce(archived_at,statement_timestamp()),updated_at=statement_timestamp() +WHERE id=$1 AND state='active' +` + +type ArchiveExecutionAgentParams struct { + ID uuid.UUID +} + +func (q *Queries) ArchiveExecutionAgent(ctx context.Context, db DBTX, arg ArchiveExecutionAgentParams) (int64, error) { + result, err := db.Exec(ctx, archiveExecutionAgent, arg.ID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const clearExecutionIntegrationTargets = `-- name: ClearExecutionIntegrationTargets :exec +UPDATE agents a SET interaction_target_id=NULL,interaction_handler_key=NULL,updated_at=statement_timestamp() +FROM integration_targets t WHERE a.project_id=$1 AND a.id=ANY($2::uuid[]) + AND t.project_id=a.project_id AND t.agent_id=a.id AND t.id=a.interaction_target_id AND t.integration_id=$3 +` + +type ClearExecutionIntegrationTargetsParams struct { + ProjectID uuid.UUID + AgentIds []uuid.UUID + IntegrationID uuid.UUID +} + +// @sqlc-vet-disable integration-targets-deleted-at +func (q *Queries) ClearExecutionIntegrationTargets(ctx context.Context, db DBTX, arg ClearExecutionIntegrationTargetsParams) error { + _, err := db.Exec(ctx, clearExecutionIntegrationTargets, arg.ProjectID, arg.AgentIds, arg.IntegrationID) + return err +} + +const createExecutionAgent = `-- name: CreateExecutionAgent :one +INSERT INTO agents(id,org_id,project_id,root_agent_id,parent_agent_id,subagent_key,state,name,agent_profile_id, + current_config_id,idempotency_key,archive_after_idle_minutes,created_at,updated_at) + SELECT $1,p.org_id,p.id,$2,$3,$4,'active',$5, + $6,$7,$8,$9, + statement_timestamp(),statement_timestamp() + FROM projects p JOIN orgs o ON o.id=p.org_id WHERE p.id=$10 AND p.deleted_at IS NULL AND o.deleted_at IS NULL + RETURNING id +` + +type CreateExecutionAgentParams struct { + ID uuid.UUID + RootID uuid.UUID + ParentID *uuid.UUID + SubagentKey string + Name string + ProfileID *uuid.UUID + ConfigID uuid.UUID + IdempotencyKey *string + ArchiveAfterIdleMinutes *int32 + ProjectID uuid.UUID +} + +func (q *Queries) CreateExecutionAgent(ctx context.Context, db DBTX, arg CreateExecutionAgentParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, createExecutionAgent, + arg.ID, + arg.RootID, + arg.ParentID, + arg.SubagentKey, + arg.Name, + arg.ProfileID, + arg.ConfigID, + arg.IdempotencyKey, + arg.ArchiveAfterIdleMinutes, + arg.ProjectID, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const listExecutionIntegrationTargetAgents = `-- name: ListExecutionIntegrationTargetAgents :many +SELECT a.id FROM agents a JOIN integration_targets t ON t.project_id=a.project_id AND t.agent_id=a.id AND t.id=a.interaction_target_id +WHERE a.project_id=$1 AND t.integration_id=$2 ORDER BY a.id +` + +type ListExecutionIntegrationTargetAgentsParams struct { + ProjectID uuid.UUID + IntegrationID uuid.UUID +} + +// @sqlc-vet-disable integration-targets-deleted-at +func (q *Queries) ListExecutionIntegrationTargetAgents(ctx context.Context, db DBTX, arg ListExecutionIntegrationTargetAgentsParams) ([]uuid.UUID, error) { + rows, err := db.Query(ctx, listExecutionIntegrationTargetAgents, arg.ProjectID, arg.IntegrationID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []uuid.UUID{} + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const readExecutionAgentIdentity = `-- name: ReadExecutionAgentIdentity :one +SELECT org_id,parent_agent_id,root_agent_id,state,name,subagent_key FROM agents +WHERE project_id=$1 AND id=$2 +` + +type ReadExecutionAgentIdentityParams struct { + ProjectID uuid.UUID + ID uuid.UUID +} + +type ReadExecutionAgentIdentityRow struct { + OrgID uuid.UUID + ParentAgentID *uuid.UUID + RootAgentID uuid.UUID + State string + Name string + SubagentKey string +} + +func (q *Queries) ReadExecutionAgentIdentity(ctx context.Context, db DBTX, arg ReadExecutionAgentIdentityParams) (ReadExecutionAgentIdentityRow, error) { + row := db.QueryRow(ctx, readExecutionAgentIdentity, arg.ProjectID, arg.ID) + var i ReadExecutionAgentIdentityRow + err := row.Scan( + &i.OrgID, + &i.ParentAgentID, + &i.RootAgentID, + &i.State, + &i.Name, + &i.SubagentKey, + ) + return i, err +} + +const writeExecutionInteractionTarget = `-- name: WriteExecutionInteractionTarget :execrows +UPDATE agents a SET interaction_target_id=$1,interaction_handler_key=$2, + interaction_auto_select=$3,updated_at=statement_timestamp() +WHERE a.id=$4 AND (($1::uuid IS NULL AND $2::text IS NULL) + OR ($2::text<>'' AND EXISTS(SELECT 1 FROM integration_targets t JOIN integrations i + ON i.project_id=t.project_id AND i.id=t.integration_id WHERE t.project_id=a.project_id AND t.agent_id=a.id + AND t.id=$1 AND t.deleted_at IS NULL AND i.deleted_at IS NULL AND i.state='active'))) +` + +type WriteExecutionInteractionTargetParams struct { + TargetID *uuid.UUID + HandlerKey *string + AutoSelect bool + ID uuid.UUID +} + +func (q *Queries) WriteExecutionInteractionTarget(ctx context.Context, db DBTX, arg WriteExecutionInteractionTargetParams) (int64, error) { + result, err := db.Exec(ctx, writeExecutionInteractionTarget, + arg.TargetID, + arg.HandlerKey, + arg.AutoSelect, + arg.ID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_events.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_events.sql.go new file mode 100644 index 000000000..b4c14809e --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_events.sql.go @@ -0,0 +1,249 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_events.sql + +package executiondb + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" +) + +const advanceExecutionTurn = `-- name: AdvanceExecutionTurn :execrows +UPDATE agent_turns SET latest_event_id=$1, + latest_semantic_event_id=CASE WHEN $2::boolean THEN $1 ELSE latest_semantic_event_id END +WHERE agent_id=$3 AND id=$4 +` + +type AdvanceExecutionTurnParams struct { + EventID uuid.UUID + Semantic bool + AgentID uuid.UUID + TurnID uuid.UUID +} + +func (q *Queries) AdvanceExecutionTurn(ctx context.Context, db DBTX, arg AdvanceExecutionTurnParams) (int64, error) { + result, err := db.Exec(ctx, advanceExecutionTurn, + arg.EventID, + arg.Semantic, + arg.AgentID, + arg.TurnID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const appendExecutionEvent = `-- name: AppendExecutionEvent :one +WITH payload AS MATERIALIZED ( + SELECT agent_payload_created_at($11,$2,CASE WHEN $8::uuid IS NOT NULL THEN 'tool_call_result' ELSE $4 END, + coalesce($5::uuid,$6::uuid,$7::uuid,$8::uuid),true) AS created_at +), allocated AS ( + UPDATE agents SET next_event_sequence=next_event_sequence+1 WHERE id=$2 + RETURNING next_event_sequence-1 AS sequence +) +INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,agent_input_id,model_output_id, + context_checkpoint_id,tool_call_result_id,is_opening_event,created_at,idempotency_key) +SELECT $1,$2,$3,sequence,$4,$5, + $6,$7,$8,$9,payload.created_at, + $10 FROM allocated CROSS JOIN payload +RETURNING id,sequence,created_at +` + +type AppendExecutionEventParams struct { + ID uuid.UUID + AgentID uuid.UUID + TurnID *uuid.UUID + Kind string + InputID *uuid.UUID + OutputID *uuid.UUID + CheckpointID *uuid.UUID + ResultID *uuid.UUID + Opening bool + IdempotencyKey *string + ProjectID uuid.UUID +} + +type AppendExecutionEventRow struct { + ID uuid.UUID + Sequence int64 + CreatedAt time.Time +} + +func (q *Queries) AppendExecutionEvent(ctx context.Context, db DBTX, arg AppendExecutionEventParams) (AppendExecutionEventRow, error) { + row := db.QueryRow(ctx, appendExecutionEvent, + arg.ID, + arg.AgentID, + arg.TurnID, + arg.Kind, + arg.InputID, + arg.OutputID, + arg.CheckpointID, + arg.ResultID, + arg.Opening, + arg.IdempotencyKey, + arg.ProjectID, + ) + var i AppendExecutionEventRow + err := row.Scan(&i.ID, &i.Sequence, &i.CreatedAt) + return i, err +} + +const findExecutionInputEvent = `-- name: FindExecutionInputEvent :one +SELECT e.id,e.turn_id,e.sequence,e.created_at FROM agent_events e WHERE e.agent_id=$1 + AND e.agent_input_id=$2 +` + +type FindExecutionInputEventParams struct { + AgentID uuid.UUID + InputID *uuid.UUID +} + +type FindExecutionInputEventRow struct { + ID uuid.UUID + TurnID uuid.UUID + Sequence int64 + CreatedAt time.Time +} + +func (q *Queries) FindExecutionInputEvent(ctx context.Context, db DBTX, arg FindExecutionInputEventParams) (FindExecutionInputEventRow, error) { + row := db.QueryRow(ctx, findExecutionInputEvent, arg.AgentID, arg.InputID) + var i FindExecutionInputEventRow + err := row.Scan( + &i.ID, + &i.TurnID, + &i.Sequence, + &i.CreatedAt, + ) + return i, err +} + +const insertExecutionContent = `-- name: InsertExecutionContent :exec +WITH payload AS MATERIALIZED ( + SELECT agent_payload_created_at($14,$1,$2, + coalesce($3::uuid,$4::uuid,$5::uuid),false) AS created_at +) +INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,owner_model_output_id,owner_tool_call_result_id, + ordinal,block_kind,text_content,structured_data,artifact_id,tool_call_id,exclude_from_model_context,metadata,created_at) +SELECT $1,$2,$3,$4,$5, + $6,$7,$8,$9::jsonb,$10,$11, + $12,$13::jsonb,payload.created_at FROM payload +` + +type InsertExecutionContentParams struct { + AgentID uuid.UUID + OwnerKind string + InputID *uuid.UUID + OutputID *uuid.UUID + ResultID *uuid.UUID + Ordinal int32 + Kind string + Text *string + Data *json.RawMessage + ArtifactID *uuid.UUID + ToolID *uuid.UUID + Exclude bool + Metadata json.RawMessage + ProjectID uuid.UUID +} + +func (q *Queries) InsertExecutionContent(ctx context.Context, db DBTX, arg InsertExecutionContentParams) error { + _, err := db.Exec(ctx, insertExecutionContent, + arg.AgentID, + arg.OwnerKind, + arg.InputID, + arg.OutputID, + arg.ResultID, + arg.Ordinal, + arg.Kind, + arg.Text, + arg.Data, + arg.ArtifactID, + arg.ToolID, + arg.Exclude, + arg.Metadata, + arg.ProjectID, + ) + return err +} + +const openExecutionTurn = `-- name: OpenExecutionTurn :exec +INSERT INTO agent_turns(id,agent_id,turn_sequence,latest_event_id,latest_semantic_event_id) +SELECT $1,$2,coalesce((SELECT max(turn_sequence)+1 FROM agent_turns + WHERE agent_id=$2),1),$3,$3 +` + +type OpenExecutionTurnParams struct { + ID uuid.UUID + AgentID uuid.UUID + EventID uuid.UUID +} + +func (q *Queries) OpenExecutionTurn(ctx context.Context, db DBTX, arg OpenExecutionTurnParams) error { + _, err := db.Exec(ctx, openExecutionTurn, arg.ID, arg.AgentID, arg.EventID) + return err +} + +const readExecutionContent = `-- name: ReadExecutionContent :many +SELECT ordinal,block_kind,text_content,structured_data,artifact_id,tool_call_id,exclude_from_model_context,metadata +FROM content_blocks WHERE agent_id=$1 AND + (owner_agent_input_id=$2 OR owner_model_output_id=$3 + OR owner_tool_call_result_id=$4) ORDER BY ordinal +` + +type ReadExecutionContentParams struct { + AgentID uuid.UUID + InputID *uuid.UUID + OutputID *uuid.UUID + ResultID *uuid.UUID +} + +type ReadExecutionContentRow struct { + Ordinal int32 + BlockKind string + TextContent *string + StructuredData *json.RawMessage + ArtifactID *uuid.UUID + ToolCallID *uuid.UUID + ExcludeFromModelContext bool + Metadata json.RawMessage +} + +func (q *Queries) ReadExecutionContent(ctx context.Context, db DBTX, arg ReadExecutionContentParams) ([]ReadExecutionContentRow, error) { + rows, err := db.Query(ctx, readExecutionContent, + arg.AgentID, + arg.InputID, + arg.OutputID, + arg.ResultID, + ) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ReadExecutionContentRow{} + for rows.Next() { + var i ReadExecutionContentRow + if err := rows.Scan( + &i.Ordinal, + &i.BlockKind, + &i.TextContent, + &i.StructuredData, + &i.ArtifactID, + &i.ToolCallID, + &i.ExcludeFromModelContext, + &i.Metadata, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go new file mode 100644 index 000000000..a0e90e27b --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go @@ -0,0 +1,242 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_inputs.sql + +package executiondb + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" +) + +const activateExecutionConfig = `-- name: ActivateExecutionConfig :execrows +UPDATE agents a SET current_config_id=c.id,updated_at=statement_timestamp() +FROM agent_configs c WHERE a.id=$1 AND a.project_id=$2 + AND c.project_id=a.project_id AND c.id=$3 +` + +type ActivateExecutionConfigParams struct { + AgentID uuid.UUID + ProjectID uuid.UUID + ConfigID uuid.UUID +} + +func (q *Queries) ActivateExecutionConfig(ctx context.Context, db DBTX, arg ActivateExecutionConfigParams) (int64, error) { + result, err := db.Exec(ctx, activateExecutionConfig, arg.AgentID, arg.ProjectID, arg.ConfigID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const changeExecutionInput = `-- name: ChangeExecutionInput :execrows +UPDATE agent_inputs SET delivery_mode=$1,input_rank=$2, + state=CASE WHEN $3::boolean THEN 'canceled' ELSE state END, + canceled_at=CASE WHEN $3::boolean THEN statement_timestamp() ELSE NULL END +WHERE agent_id=$4 AND id=$5 AND state='received' AND input_kind='content' + AND (delivery_mode,input_rank,state) IS DISTINCT FROM + ($1,$2,CASE WHEN $3::boolean THEN 'canceled' ELSE state END) +` + +type ChangeExecutionInputParams struct { + Mode string + Rank int64 + Cancel bool + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) ChangeExecutionInput(ctx context.Context, db DBTX, arg ChangeExecutionInputParams) (int64, error) { + result, err := db.Exec(ctx, changeExecutionInput, + arg.Mode, + arg.Rank, + arg.Cancel, + arg.AgentID, + arg.ID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const findExecutionInput = `-- name: FindExecutionInput :one +SELECT id,state,input_kind,delivery_mode,actor_id,integration_target_id,agent_config_id, + metadata,queued_at,admitted_event_id,input_rank +FROM agent_inputs WHERE agent_id=$1 AND + (id=$2 OR (idempotency_scope=$3 AND input_idempotency_key=$4)) +` + +type FindExecutionInputParams struct { + AgentID uuid.UUID + ID uuid.UUID + Scope *string + Key *string +} + +type FindExecutionInputRow struct { + ID uuid.UUID + State string + InputKind string + DeliveryMode string + ActorID *uuid.UUID + IntegrationTargetID *uuid.UUID + AgentConfigID *uuid.UUID + Metadata json.RawMessage + QueuedAt time.Time + AdmittedEventID *uuid.UUID + InputRank int64 +} + +func (q *Queries) FindExecutionInput(ctx context.Context, db DBTX, arg FindExecutionInputParams) (FindExecutionInputRow, error) { + row := db.QueryRow(ctx, findExecutionInput, + arg.AgentID, + arg.ID, + arg.Scope, + arg.Key, + ) + var i FindExecutionInputRow + err := row.Scan( + &i.ID, + &i.State, + &i.InputKind, + &i.DeliveryMode, + &i.ActorID, + &i.IntegrationTargetID, + &i.AgentConfigID, + &i.Metadata, + &i.QueuedAt, + &i.AdmittedEventID, + &i.InputRank, + ) + return i, err +} + +const listExecutionBacklog = `-- name: ListExecutionBacklog :many +SELECT id,delivery_mode,input_rank,queued_at,integration_target_id,actor_id FROM agent_inputs +WHERE agent_id=$1 AND input_kind='content' AND state='received' +ORDER BY input_rank,queued_at,id +` + +type ListExecutionBacklogParams struct { + AgentID uuid.UUID +} + +type ListExecutionBacklogRow struct { + ID uuid.UUID + DeliveryMode string + InputRank int64 + QueuedAt time.Time + IntegrationTargetID *uuid.UUID + ActorID *uuid.UUID +} + +func (q *Queries) ListExecutionBacklog(ctx context.Context, db DBTX, arg ListExecutionBacklogParams) ([]ListExecutionBacklogRow, error) { + rows, err := db.Query(ctx, listExecutionBacklog, arg.AgentID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ListExecutionBacklogRow{} + for rows.Next() { + var i ListExecutionBacklogRow + if err := rows.Scan( + &i.ID, + &i.DeliveryMode, + &i.InputRank, + &i.QueuedAt, + &i.IntegrationTargetID, + &i.ActorID, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const receiveExecutionInput = `-- name: ReceiveExecutionInput :one +INSERT INTO agent_inputs(id,project_id,agent_id,state,input_rank,actor_id,input_kind,delivery_mode, + integration_target_id,agent_config_id,idempotency_scope,input_idempotency_key,queued_at,metadata) +SELECT $1,a.project_id,a.id,'received',coalesce((SELECT max(i.input_rank)+1024 + FROM agent_inputs i WHERE i.agent_id=a.id AND i.state='received' AND i.delivery_mode=$2),1024), + $3,$4,$2,$5,$6, + $7,$8,statement_timestamp(),$9::jsonb +FROM agents a WHERE a.id=$10 AND a.project_id=$11 AND a.state='active' +RETURNING id,queued_at +` + +type ReceiveExecutionInputParams struct { + ID uuid.UUID + Mode string + ActorID *uuid.UUID + Kind string + TargetID *uuid.UUID + ConfigID *uuid.UUID + Scope *string + Key *string + Metadata json.RawMessage + AgentID uuid.UUID + ProjectID uuid.UUID +} + +type ReceiveExecutionInputRow struct { + ID uuid.UUID + QueuedAt time.Time +} + +func (q *Queries) ReceiveExecutionInput(ctx context.Context, db DBTX, arg ReceiveExecutionInputParams) (ReceiveExecutionInputRow, error) { + row := db.QueryRow(ctx, receiveExecutionInput, + arg.ID, + arg.Mode, + arg.ActorID, + arg.Kind, + arg.TargetID, + arg.ConfigID, + arg.Scope, + arg.Key, + arg.Metadata, + arg.AgentID, + arg.ProjectID, + ) + var i ReceiveExecutionInputRow + err := row.Scan(&i.ID, &i.QueuedAt) + return i, err +} + +const resolveExecutionInput = `-- name: ResolveExecutionInput :execrows +UPDATE agent_inputs i SET state='resolved',admitted_event_id=$1, + admitted_at=e.created_at,resolved_at=e.created_at, + opening_input_ids=$2::uuid[],opening_event_sequence=$3::bigint +FROM agent_events e WHERE i.agent_id=$4 AND i.id=$5 AND i.state='received' + AND e.agent_id=i.agent_id AND e.id=$1 AND e.agent_input_id=i.id +` + +type ResolveExecutionInputParams struct { + EventID *uuid.UUID + OpeningIds []uuid.UUID + OpeningSequence *int64 + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) ResolveExecutionInput(ctx context.Context, db DBTX, arg ResolveExecutionInputParams) (int64, error) { + result, err := db.Exec(ctx, resolveExecutionInput, + arg.EventID, + arg.OpeningIds, + arg.OpeningSequence, + arg.AgentID, + arg.ID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_interactions.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_interactions.sql.go new file mode 100644 index 000000000..e28c1288b --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_interactions.sql.go @@ -0,0 +1,202 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_interactions.sql + +package executiondb + +import ( + "context" + "encoding/json" + + "github.com/google/uuid" +) + +const insertExecutionInteraction = `-- name: InsertExecutionInteraction :one +INSERT INTO agent_interactions(agent_id,tool_call_id,interaction_kind,state,request,destination,created_at) +VALUES($1,$2,$3,'open',$4::jsonb, + $5::jsonb,statement_timestamp()) RETURNING id +` + +type InsertExecutionInteractionParams struct { + AgentID uuid.UUID + ToolID uuid.UUID + Kind string + Request json.RawMessage + Destination *json.RawMessage +} + +func (q *Queries) InsertExecutionInteraction(ctx context.Context, db DBTX, arg InsertExecutionInteractionParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, insertExecutionInteraction, + arg.AgentID, + arg.ToolID, + arg.Kind, + arg.Request, + arg.Destination, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const insertExecutionResponse = `-- name: InsertExecutionResponse :one +INSERT INTO agent_inputs(project_id,agent_id,state,actor_id,input_kind,delivery_mode,target_interaction_id, + idempotency_scope,input_idempotency_key,queued_at,metadata) +VALUES($1,$2,'received',$3,'interaction_response','immediate', + $4,'agent_interaction_response',$4::uuid::text,statement_timestamp(),'{}'::jsonb) +RETURNING id +` + +type InsertExecutionResponseParams struct { + ProjectID uuid.UUID + AgentID uuid.UUID + ActorID *uuid.UUID + InteractionID *uuid.UUID +} + +func (q *Queries) InsertExecutionResponse(ctx context.Context, db DBTX, arg InsertExecutionResponseParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, insertExecutionResponse, + arg.ProjectID, + arg.AgentID, + arg.ActorID, + arg.InteractionID, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const listExecutionInteractions = `-- name: ListExecutionInteractions :many +SELECT i.id FROM agent_interactions i JOIN tool_calls c ON c.agent_id=i.agent_id AND c.id=i.tool_call_id +JOIN model_outputs o ON o.agent_id=c.agent_id AND o.id=c.model_output_id +JOIN model_call_contexts m ON m.agent_id=o.agent_id AND m.id=o.model_call_context_id +WHERE i.agent_id=$1 AND i.state='open' AND m.turn_id=$2 +ORDER BY i.id +` + +type ListExecutionInteractionsParams struct { + AgentID uuid.UUID + TurnID uuid.UUID +} + +func (q *Queries) ListExecutionInteractions(ctx context.Context, db DBTX, arg ListExecutionInteractionsParams) ([]uuid.UUID, error) { + rows, err := db.Query(ctx, listExecutionInteractions, arg.AgentID, arg.TurnID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []uuid.UUID{} + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const readExecutionInteraction = `-- name: ReadExecutionInteraction :one +SELECT i.id,i.tool_call_id,i.interaction_kind,i.state,i.request,i.resolution,i.resolved_by_input_id, + i.destination,i.presentation_receipt +FROM agent_interactions i WHERE i.agent_id=$1 AND + (i.id=$2 OR (i.tool_call_id=$3 AND i.interaction_kind=$4)) +` + +type ReadExecutionInteractionParams struct { + AgentID uuid.UUID + ID uuid.UUID + ToolID *uuid.UUID + Kind *string +} + +type ReadExecutionInteractionRow struct { + ID uuid.UUID + ToolCallID uuid.UUID + InteractionKind string + State string + Request json.RawMessage + Resolution json.RawMessage + ResolvedByInputID *uuid.UUID + Destination *json.RawMessage + PresentationReceipt *json.RawMessage +} + +func (q *Queries) ReadExecutionInteraction(ctx context.Context, db DBTX, arg ReadExecutionInteractionParams) (ReadExecutionInteractionRow, error) { + row := db.QueryRow(ctx, readExecutionInteraction, + arg.AgentID, + arg.ID, + arg.ToolID, + arg.Kind, + ) + var i ReadExecutionInteractionRow + err := row.Scan( + &i.ID, + &i.ToolCallID, + &i.InteractionKind, + &i.State, + &i.Request, + &i.Resolution, + &i.ResolvedByInputID, + &i.Destination, + &i.PresentationReceipt, + ) + return i, err +} + +const resolveExecutionInteraction = `-- name: ResolveExecutionInteraction :execrows +UPDATE agent_interactions SET state=$1,resolution=$2::jsonb, + resolved_by_input_id=$3,resolved_at=statement_timestamp() +WHERE agent_id=$4 AND id=$5 AND state='open' +` + +type ResolveExecutionInteractionParams struct { + State string + Resolution json.RawMessage + InputID *uuid.UUID + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) ResolveExecutionInteraction(ctx context.Context, db DBTX, arg ResolveExecutionInteractionParams) (int64, error) { + result, err := db.Exec(ctx, resolveExecutionInteraction, + arg.State, + arg.Resolution, + arg.InputID, + arg.AgentID, + arg.ID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const writeExecutionPresentation = `-- name: WriteExecutionPresentation :execrows +UPDATE agent_interactions SET presentation_receipt=$1::jsonb +WHERE agent_id=$2 AND id=$3 AND destination=$4::jsonb + AND presentation_receipt IS NULL +` + +type WriteExecutionPresentationParams struct { + Receipt json.RawMessage + AgentID uuid.UUID + ID uuid.UUID + Destination json.RawMessage +} + +func (q *Queries) WriteExecutionPresentation(ctx context.Context, db DBTX, arg WriteExecutionPresentationParams) (int64, error) { + result, err := db.Exec(ctx, writeExecutionPresentation, + arg.Receipt, + arg.AgentID, + arg.ID, + arg.Destination, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_models.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_models.sql.go new file mode 100644 index 000000000..d8320109c --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_models.sql.go @@ -0,0 +1,349 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_models.sql + +package executiondb + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" +) + +const createExecutionAttempt = `-- name: CreateExecutionAttempt :one +INSERT INTO model_call_contexts(org_id,project_id,agent_id,id,turn_id,operation_kind,attempt_number, + agent_config_id,configured_model_revision_id,input_event_sequence,source_event_sequence_end,runtime_lock_id, + opening_input_ids,opening_event_sequence,state,created_at) +SELECT a.org_id,a.project_id,a.id,$1,$2,$3,$4, + $5,$6,$7,$8,$9, + $10::uuid[],$11,'started',statement_timestamp() +FROM agents a WHERE a.id=$12 AND a.project_id=$13 + AND EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id AND r.id=$9 + AND r.cancel_requested_at IS NULL AND r.lease_expires_at>statement_timestamp()) +RETURNING id +` + +type CreateExecutionAttemptParams struct { + ID uuid.UUID + TurnID uuid.UUID + Operation string + Attempt int32 + ConfigID uuid.UUID + RevisionID uuid.UUID + Watermark int64 + SourceEnd *int64 + RuntimeID uuid.UUID + OpeningIds []uuid.UUID + OpeningSequence int64 + AgentID uuid.UUID + ProjectID uuid.UUID +} + +func (q *Queries) CreateExecutionAttempt(ctx context.Context, db DBTX, arg CreateExecutionAttemptParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, createExecutionAttempt, + arg.ID, + arg.TurnID, + arg.Operation, + arg.Attempt, + arg.ConfigID, + arg.RevisionID, + arg.Watermark, + arg.SourceEnd, + arg.RuntimeID, + arg.OpeningIds, + arg.OpeningSequence, + arg.AgentID, + arg.ProjectID, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const executionCompactionBoundary = `-- name: ExecutionCompactionBoundary :one +SELECT coalesce((SELECT summarized_through_event_sequence FROM context_checkpoints + WHERE agent_id=$1 AND summarized_through_event_sequence<=$2 + ORDER BY summarized_through_event_sequence DESC LIMIT 1),0)::bigint AS summarized_through +` + +type ExecutionCompactionBoundaryParams struct { + AgentID uuid.UUID + Watermark int64 +} + +func (q *Queries) ExecutionCompactionBoundary(ctx context.Context, db DBTX, arg ExecutionCompactionBoundaryParams) (int64, error) { + row := db.QueryRow(ctx, executionCompactionBoundary, arg.AgentID, arg.Watermark) + var summarized_through int64 + err := row.Scan(&summarized_through) + return summarized_through, err +} + +const executionModelRevision = `-- name: ExecutionModelRevision :one +SELECT m.current_revision_id,CASE WHEN p.management_kind='cluster' THEN coalesce(w.new_managed_work_allowed,true) + ELSE true END::boolean AS allowed +FROM agent_configs c JOIN configured_models m ON m.org_id=c.org_id AND m.id=c.configured_model_id +JOIN model_provider_configs p ON p.org_id=m.org_id AND p.id=m.model_provider_config_id +LEFT JOIN org_managed_work_admission w ON w.org_id=c.org_id +WHERE c.project_id=$1 AND c.id=$2 +` + +type ExecutionModelRevisionParams struct { + ProjectID uuid.UUID + ConfigID uuid.UUID +} + +type ExecutionModelRevisionRow struct { + CurrentRevisionID uuid.UUID + Allowed bool +} + +// @sqlc-vet-disable configured-models-deleted-at +// @sqlc-vet-disable model-provider-configs-deleted-at +func (q *Queries) ExecutionModelRevision(ctx context.Context, db DBTX, arg ExecutionModelRevisionParams) (ExecutionModelRevisionRow, error) { + row := db.QueryRow(ctx, executionModelRevision, arg.ProjectID, arg.ConfigID) + var i ExecutionModelRevisionRow + err := row.Scan(&i.CurrentRevisionID, &i.Allowed) + return i, err +} + +const executionPreparationBoundary = `-- name: ExecutionPreparationBoundary :one +SELECT a.current_config_id,(a.next_event_sequence-1)::bigint AS watermark,statement_timestamp()::timestamptz AS database_now, + EXISTS(SELECT 1 FROM model_call_contexts c WHERE c.agent_id=a.id AND c.state='started') AS started +FROM agents a WHERE a.id=$1 AND a.state='active' +` + +type ExecutionPreparationBoundaryParams struct { + AgentID uuid.UUID +} + +type ExecutionPreparationBoundaryRow struct { + CurrentConfigID uuid.UUID + Watermark int64 + DatabaseNow time.Time + Started bool +} + +func (q *Queries) ExecutionPreparationBoundary(ctx context.Context, db DBTX, arg ExecutionPreparationBoundaryParams) (ExecutionPreparationBoundaryRow, error) { + row := db.QueryRow(ctx, executionPreparationBoundary, arg.AgentID) + var i ExecutionPreparationBoundaryRow + err := row.Scan( + &i.CurrentConfigID, + &i.Watermark, + &i.DatabaseNow, + &i.Started, + ) + return i, err +} + +const fenceExecutionRuntime = `-- name: FenceExecutionRuntime :one +SELECT r.id FROM agent_runtime_locks r JOIN agents a ON a.id=r.agent_id +WHERE a.project_id=$1 AND a.id=$2 AND a.state='active' + AND r.id=$3 AND r.cancel_requested_at IS NULL AND r.lease_expires_at>statement_timestamp() +` + +type FenceExecutionRuntimeParams struct { + ProjectID uuid.UUID + AgentID uuid.UUID + RuntimeID uuid.UUID +} + +func (q *Queries) FenceExecutionRuntime(ctx context.Context, db DBTX, arg FenceExecutionRuntimeParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, fenceExecutionRuntime, arg.ProjectID, arg.AgentID, arg.RuntimeID) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const finishExecutionAttempt = `-- name: FinishExecutionAttempt :execrows +UPDATE model_call_contexts c SET state=$1,recovery_kind=$2, + retry_at=CASE WHEN $3::bigint IS NULL THEN NULL + ELSE statement_timestamp()+$3::bigint*interval '1 microsecond' END, + api_format=$4,api_variant=$5,provider_request_id=$6, + provider_response_id=$7,error_kind=$8,error_code=$9, + error_message=$10,error_details=$11::jsonb, + input_tokens_total=$12::integer,uncached_input_tokens=$13::integer, + cache_read_input_tokens=$14::integer,cache_write_input_tokens=$15::integer, + output_tokens_total=$16::integer,reasoning_output_tokens=$17::integer, + provider_reported_cost_usd=$18::text::numeric,provider_metadata=$19::jsonb, + completed_at=statement_timestamp() +WHERE c.agent_id=$20 AND c.id=$21 AND c.state='started' AND c.runtime_lock_id=$22 + AND EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=c.agent_id + AND r.id=$22 AND r.cancel_requested_at IS NULL AND r.lease_expires_at>statement_timestamp()) +` + +type FinishExecutionAttemptParams struct { + State string + Recovery *string + RetryMicroseconds *int64 + ApiFormat string + ApiVariant string + RequestID string + ResponseID string + ErrorKind string + ErrorCode string + ErrorMessage string + ErrorDetails json.RawMessage + InputTokens *int32 + UncachedTokens *int32 + CacheReadTokens *int32 + CacheWriteTokens *int32 + OutputTokens *int32 + ReasoningTokens *int32 + Cost *string + ProviderMetadata json.RawMessage + AgentID uuid.UUID + ID uuid.UUID + RuntimeID uuid.UUID +} + +func (q *Queries) FinishExecutionAttempt(ctx context.Context, db DBTX, arg FinishExecutionAttemptParams) (int64, error) { + result, err := db.Exec(ctx, finishExecutionAttempt, + arg.State, + arg.Recovery, + arg.RetryMicroseconds, + arg.ApiFormat, + arg.ApiVariant, + arg.RequestID, + arg.ResponseID, + arg.ErrorKind, + arg.ErrorCode, + arg.ErrorMessage, + arg.ErrorDetails, + arg.InputTokens, + arg.UncachedTokens, + arg.CacheReadTokens, + arg.CacheWriteTokens, + arg.OutputTokens, + arg.ReasoningTokens, + arg.Cost, + arg.ProviderMetadata, + arg.AgentID, + arg.ID, + arg.RuntimeID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const latestExecutionAttempt = `-- name: LatestExecutionAttempt :one +SELECT id FROM model_call_contexts WHERE agent_id=$1 AND operation_kind=$2 + AND input_event_sequence=$3 + AND source_event_sequence_end IS NOT DISTINCT FROM $4::bigint +ORDER BY attempt_number DESC LIMIT 1 +` + +type LatestExecutionAttemptParams struct { + AgentID uuid.UUID + Operation string + Watermark int64 + SourceEnd *int64 +} + +func (q *Queries) LatestExecutionAttempt(ctx context.Context, db DBTX, arg LatestExecutionAttemptParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, latestExecutionAttempt, + arg.AgentID, + arg.Operation, + arg.Watermark, + arg.SourceEnd, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const readExecutionAttempt = `-- name: ReadExecutionAttempt :one +SELECT c.id,c.agent_id,c.turn_id,c.operation_kind,c.attempt_number,c.input_event_sequence,c.source_event_sequence_end, + c.state,c.recovery_kind,c.retry_at,c.opening_input_ids,c.opening_event_sequence,c.runtime_lock_id, + c.agent_config_id,c.configured_model_revision_id,c.api_format,c.api_variant,c.provider_request_id,c.provider_response_id, + c.error_kind,c.error_code,c.error_message,c.error_details,c.input_tokens_total,c.uncached_input_tokens, + c.cache_read_input_tokens,c.cache_write_input_tokens,c.output_tokens_total,c.reasoning_output_tokens, + coalesce(c.provider_reported_cost_usd::text,'')::text AS provider_reported_cost_usd,c.provider_metadata, + r.provider_model_slug +FROM model_call_contexts c JOIN configured_model_revisions r ON r.org_id=c.org_id AND r.id=c.configured_model_revision_id +WHERE c.agent_id=$1 AND c.id=$2 +` + +type ReadExecutionAttemptParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +type ReadExecutionAttemptRow struct { + ID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + OperationKind string + AttemptNumber int32 + InputEventSequence int64 + SourceEventSequenceEnd *int64 + State string + RecoveryKind *string + RetryAt *time.Time + OpeningInputIds []uuid.UUID + OpeningEventSequence int64 + RuntimeLockID uuid.UUID + AgentConfigID uuid.UUID + ConfiguredModelRevisionID uuid.UUID + ApiFormat string + ApiVariant string + ProviderRequestID string + ProviderResponseID string + ErrorKind string + ErrorCode string + ErrorMessage string + ErrorDetails json.RawMessage + InputTokensTotal *int32 + UncachedInputTokens *int32 + CacheReadInputTokens *int32 + CacheWriteInputTokens *int32 + OutputTokensTotal *int32 + ReasoningOutputTokens *int32 + ProviderReportedCostUsd string + ProviderMetadata json.RawMessage + ProviderModelSlug string +} + +func (q *Queries) ReadExecutionAttempt(ctx context.Context, db DBTX, arg ReadExecutionAttemptParams) (ReadExecutionAttemptRow, error) { + row := db.QueryRow(ctx, readExecutionAttempt, arg.AgentID, arg.ID) + var i ReadExecutionAttemptRow + err := row.Scan( + &i.ID, + &i.AgentID, + &i.TurnID, + &i.OperationKind, + &i.AttemptNumber, + &i.InputEventSequence, + &i.SourceEventSequenceEnd, + &i.State, + &i.RecoveryKind, + &i.RetryAt, + &i.OpeningInputIds, + &i.OpeningEventSequence, + &i.RuntimeLockID, + &i.AgentConfigID, + &i.ConfiguredModelRevisionID, + &i.ApiFormat, + &i.ApiVariant, + &i.ProviderRequestID, + &i.ProviderResponseID, + &i.ErrorKind, + &i.ErrorCode, + &i.ErrorMessage, + &i.ErrorDetails, + &i.InputTokensTotal, + &i.UncachedInputTokens, + &i.CacheReadInputTokens, + &i.CacheWriteInputTokens, + &i.OutputTokensTotal, + &i.ReasoningOutputTokens, + &i.ProviderReportedCostUsd, + &i.ProviderMetadata, + &i.ProviderModelSlug, + ) + return i, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_outputs.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_outputs.sql.go new file mode 100644 index 000000000..eb28d77a4 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_outputs.sql.go @@ -0,0 +1,258 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_outputs.sql + +package executiondb + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" +) + +const createExecutionCheckpoint = `-- name: CreateExecutionCheckpoint :one +INSERT INTO context_checkpoints(agent_id,summarized_through_event_sequence,producer_model_call_context_id,summary, + opening_input_ids,opening_event_sequence,created_at) +VALUES($1,$2,$3,$4, + $5::uuid[],$6,statement_timestamp()) +RETURNING id,created_at +` + +type CreateExecutionCheckpointParams struct { + AgentID uuid.UUID + SourceEnd int64 + ContextID uuid.UUID + Summary string + OpeningIds []uuid.UUID + OpeningSequence *int64 +} + +type CreateExecutionCheckpointRow struct { + ID uuid.UUID + CreatedAt time.Time +} + +func (q *Queries) CreateExecutionCheckpoint(ctx context.Context, db DBTX, arg CreateExecutionCheckpointParams) (CreateExecutionCheckpointRow, error) { + row := db.QueryRow(ctx, createExecutionCheckpoint, + arg.AgentID, + arg.SourceEnd, + arg.ContextID, + arg.Summary, + arg.OpeningIds, + arg.OpeningSequence, + ) + var i CreateExecutionCheckpointRow + err := row.Scan(&i.ID, &i.CreatedAt) + return i, err +} + +const createExecutionOutput = `-- name: CreateExecutionOutput :one +INSERT INTO model_outputs(agent_id,model_call_context_id,served_provider_model_slug,stop_reason,provider_replay,created_at) +VALUES($1,$2,$3,$4,$5::jsonb,statement_timestamp()) +RETURNING id,created_at +` + +type CreateExecutionOutputParams struct { + AgentID uuid.UUID + ContextID uuid.UUID + ServedModel string + StopReason string + Replay *json.RawMessage +} + +type CreateExecutionOutputRow struct { + ID uuid.UUID + CreatedAt time.Time +} + +func (q *Queries) CreateExecutionOutput(ctx context.Context, db DBTX, arg CreateExecutionOutputParams) (CreateExecutionOutputRow, error) { + row := db.QueryRow(ctx, createExecutionOutput, + arg.AgentID, + arg.ContextID, + arg.ServedModel, + arg.StopReason, + arg.Replay, + ) + var i CreateExecutionOutputRow + err := row.Scan(&i.ID, &i.CreatedAt) + return i, err +} + +const createExecutionToolProposal = `-- name: CreateExecutionToolProposal :exec +INSERT INTO tool_calls(id,agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) +SELECT $1,$2,o.id,$3,$4,$5::jsonb, + $6,'awaiting_authorization',o.created_at FROM model_outputs o + WHERE o.agent_id=$2 AND o.id=$7 +` + +type CreateExecutionToolProposalParams struct { + ID uuid.UUID + AgentID uuid.UUID + ProviderCallID string + Name string + Input json.RawMessage + Type string + OutputID uuid.UUID +} + +func (q *Queries) CreateExecutionToolProposal(ctx context.Context, db DBTX, arg CreateExecutionToolProposalParams) error { + _, err := db.Exec(ctx, createExecutionToolProposal, + arg.ID, + arg.AgentID, + arg.ProviderCallID, + arg.Name, + arg.Input, + arg.Type, + arg.OutputID, + ) + return err +} + +const findExecutionCheckpoint = `-- name: FindExecutionCheckpoint :one +SELECT c.id,c.summary,c.summarized_through_event_sequence,e.id AS event_id,e.turn_id,e.sequence,e.created_at +FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=$1 AND c.producer_model_call_context_id=$2 +` + +type FindExecutionCheckpointParams struct { + AgentID uuid.UUID + ContextID uuid.UUID +} + +type FindExecutionCheckpointRow struct { + ID uuid.UUID + Summary string + SummarizedThroughEventSequence int64 + EventID uuid.UUID + TurnID uuid.UUID + Sequence int64 + CreatedAt time.Time +} + +func (q *Queries) FindExecutionCheckpoint(ctx context.Context, db DBTX, arg FindExecutionCheckpointParams) (FindExecutionCheckpointRow, error) { + row := db.QueryRow(ctx, findExecutionCheckpoint, arg.AgentID, arg.ContextID) + var i FindExecutionCheckpointRow + err := row.Scan( + &i.ID, + &i.Summary, + &i.SummarizedThroughEventSequence, + &i.EventID, + &i.TurnID, + &i.Sequence, + &i.CreatedAt, + ) + return i, err +} + +const findExecutionOutput = `-- name: FindExecutionOutput :one +SELECT o.id,o.served_provider_model_slug,o.stop_reason,o.provider_replay,e.id AS event_id,e.turn_id,e.sequence,e.created_at +FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +WHERE o.agent_id=$1 AND o.model_call_context_id=$2 +` + +type FindExecutionOutputParams struct { + AgentID uuid.UUID + ContextID uuid.UUID +} + +type FindExecutionOutputRow struct { + ID uuid.UUID + ServedProviderModelSlug string + StopReason string + ProviderReplay *json.RawMessage + EventID uuid.UUID + TurnID uuid.UUID + Sequence int64 + CreatedAt time.Time +} + +func (q *Queries) FindExecutionOutput(ctx context.Context, db DBTX, arg FindExecutionOutputParams) (FindExecutionOutputRow, error) { + row := db.QueryRow(ctx, findExecutionOutput, arg.AgentID, arg.ContextID) + var i FindExecutionOutputRow + err := row.Scan( + &i.ID, + &i.ServedProviderModelSlug, + &i.StopReason, + &i.ProviderReplay, + &i.EventID, + &i.TurnID, + &i.Sequence, + &i.CreatedAt, + ) + return i, err +} + +const readExecutionProposals = `-- name: ReadExecutionProposals :many +SELECT id,provider_call_id,name,input,type FROM tool_calls +WHERE agent_id=$1 AND model_output_id=$2 ORDER BY provider_call_id +` + +type ReadExecutionProposalsParams struct { + AgentID uuid.UUID + OutputID uuid.UUID +} + +type ReadExecutionProposalsRow struct { + ID uuid.UUID + ProviderCallID string + Name string + Input json.RawMessage + Type string +} + +func (q *Queries) ReadExecutionProposals(ctx context.Context, db DBTX, arg ReadExecutionProposalsParams) ([]ReadExecutionProposalsRow, error) { + rows, err := db.Query(ctx, readExecutionProposals, arg.AgentID, arg.OutputID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ReadExecutionProposalsRow{} + for rows.Next() { + var i ReadExecutionProposalsRow + if err := rows.Scan( + &i.ID, + &i.ProviderCallID, + &i.Name, + &i.Input, + &i.Type, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const validateExecutionCheckpoint = `-- name: ValidateExecutionCheckpoint :one +SELECT (SELECT count(*) FROM agent_events e WHERE e.agent_id=$1 + AND e.sequence BETWEEN $2 AND $3)::bigint AS events, + EXISTS(SELECT 1 FROM agent_events e JOIN tool_calls c ON c.agent_id=e.agent_id AND c.model_output_id=e.model_output_id + LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id + LEFT JOIN agent_events re ON re.agent_id=r.agent_id AND re.tool_call_result_id=r.id + WHERE e.agent_id=$1 AND e.sequence BETWEEN $2 AND $3 + AND (re.id IS NULL OR re.sequence>$3)) AS open_authorities +` + +type ValidateExecutionCheckpointParams struct { + AgentID uuid.UUID + StartSequence int64 + EndSequence int64 +} + +type ValidateExecutionCheckpointRow struct { + Events int64 + OpenAuthorities bool +} + +func (q *Queries) ValidateExecutionCheckpoint(ctx context.Context, db DBTX, arg ValidateExecutionCheckpointParams) (ValidateExecutionCheckpointRow, error) { + row := db.QueryRow(ctx, validateExecutionCheckpoint, arg.AgentID, arg.StartSequence, arg.EndSequence) + var i ValidateExecutionCheckpointRow + err := row.Scan(&i.Events, &i.OpenAuthorities) + return i, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_processes.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_processes.sql.go new file mode 100644 index 000000000..aadb953bd --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_processes.sql.go @@ -0,0 +1,108 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_processes.sql + +package executiondb + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/processcmd" +) + +const readExecutionActionOwner = `-- name: ReadExecutionActionOwner :one +SELECT a.id,a.process_id,a.tool_call_id,a.action_kind,a.seq,a.payload,a.state,a.state_reason_code,a.state_reason_message, + EXISTS(SELECT 1 FROM process_actions earlier WHERE earlier.agent_id=a.agent_id AND earlier.process_id=a.process_id + AND earlier.seq'completed' + AND ($2::uuid IS NULL OR m.turn_id=$2) + AND ($3::uuid IS NULL OR (c.state='running' AND c.runtime_lock_id=$3)) ORDER BY c.id +` + +type ListExecutionUnfinishedToolsParams struct { + AgentID uuid.UUID + TurnID *uuid.UUID + RuntimeID *uuid.UUID +} + +func (q *Queries) ListExecutionUnfinishedTools(ctx context.Context, db DBTX, arg ListExecutionUnfinishedToolsParams) ([]uuid.UUID, error) { + rows, err := db.Query(ctx, listExecutionUnfinishedTools, arg.AgentID, arg.TurnID, arg.RuntimeID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []uuid.UUID{} + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const lockExecutionRuntime = `-- name: LockExecutionRuntime :one +SELECT r.id FROM agent_runtime_locks r WHERE r.agent_id=$1 AND r.id=$2 FOR UPDATE +` + +type LockExecutionRuntimeParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) LockExecutionRuntime(ctx context.Context, db DBTX, arg LockExecutionRuntimeParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, lockExecutionRuntime, arg.AgentID, arg.ID) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const readExecutionRuntime = `-- name: ReadExecutionRuntime :one +SELECT id,worker_process_id,cancel_requested_at,lease_expires_at,statement_timestamp()::timestamptz AS database_now +FROM agent_runtime_locks WHERE agent_id=$1 AND ($2::uuid IS NULL OR id=$2) +` + +type ReadExecutionRuntimeParams struct { + AgentID uuid.UUID + ID *uuid.UUID +} + +type ReadExecutionRuntimeRow struct { + ID uuid.UUID + WorkerProcessID uuid.UUID + CancelRequestedAt *time.Time + LeaseExpiresAt time.Time + DatabaseNow time.Time +} + +func (q *Queries) ReadExecutionRuntime(ctx context.Context, db DBTX, arg ReadExecutionRuntimeParams) (ReadExecutionRuntimeRow, error) { + row := db.QueryRow(ctx, readExecutionRuntime, arg.AgentID, arg.ID) + var i ReadExecutionRuntimeRow + err := row.Scan( + &i.ID, + &i.WorkerProcessID, + &i.CancelRequestedAt, + &i.LeaseExpiresAt, + &i.DatabaseNow, + ) + return i, err +} + +const renewExecutionRuntime = `-- name: RenewExecutionRuntime :one +UPDATE agent_runtime_locks r SET renewed_at=greatest(r.renewed_at,statement_timestamp()), + lease_expires_at=greatest(r.lease_expires_at,greatest(r.renewed_at,statement_timestamp())+$1::bigint*interval '1 microsecond') +WHERE r.agent_id=$2 AND r.id=$3 + AND EXISTS(SELECT 1 FROM agents a WHERE a.id=r.agent_id AND a.project_id=$4) +RETURNING r.lease_expires_at +` + +type RenewExecutionRuntimeParams struct { + LeaseMicroseconds int64 + AgentID uuid.UUID + ID uuid.UUID + ProjectID uuid.UUID +} + +func (q *Queries) RenewExecutionRuntime(ctx context.Context, db DBTX, arg RenewExecutionRuntimeParams) (time.Time, error) { + row := db.QueryRow(ctx, renewExecutionRuntime, + arg.LeaseMicroseconds, + arg.AgentID, + arg.ID, + arg.ProjectID, + ) + var lease_expires_at time.Time + err := row.Scan(&lease_expires_at) + return lease_expires_at, err +} + +const requestExecutionCancel = `-- name: RequestExecutionCancel :one +UPDATE agent_runtime_locks SET cancel_requested_at=coalesce(cancel_requested_at,statement_timestamp()) +WHERE agent_id=$1 RETURNING id,worker_process_id +` + +type RequestExecutionCancelParams struct { + AgentID uuid.UUID +} + +type RequestExecutionCancelRow struct { + ID uuid.UUID + WorkerProcessID uuid.UUID +} + +func (q *Queries) RequestExecutionCancel(ctx context.Context, db DBTX, arg RequestExecutionCancelParams) (RequestExecutionCancelRow, error) { + row := db.QueryRow(ctx, requestExecutionCancel, arg.AgentID) + var i RequestExecutionCancelRow + err := row.Scan(&i.ID, &i.WorkerProcessID) + return i, err +} + +const teardownExecutionContext = `-- name: TeardownExecutionContext :execrows +UPDATE model_call_contexts SET state=$1,recovery_kind=$2, + retry_at=CASE WHEN $3::bigint IS NULL THEN NULL + ELSE statement_timestamp()+$3::bigint*interval '1 microsecond' END, + error_kind=$4,error_code=$5,error_message=$6, + error_details=$7::jsonb,completed_at=statement_timestamp() +WHERE agent_id=$8 AND id=$9 AND runtime_lock_id=$10 AND state='started' +` + +type TeardownExecutionContextParams struct { + State string + Recovery *string + RetryMicroseconds *int64 + ErrorKind string + ErrorCode string + ErrorMessage string + ErrorDetails json.RawMessage + AgentID uuid.UUID + ID uuid.UUID + RuntimeID uuid.UUID +} + +func (q *Queries) TeardownExecutionContext(ctx context.Context, db DBTX, arg TeardownExecutionContextParams) (int64, error) { + result, err := db.Exec(ctx, teardownExecutionContext, + arg.State, + arg.Recovery, + arg.RetryMicroseconds, + arg.ErrorKind, + arg.ErrorCode, + arg.ErrorMessage, + arg.ErrorDetails, + arg.AgentID, + arg.ID, + arg.RuntimeID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const tryLockExecutionRuntime = `-- name: TryLockExecutionRuntime :one +SELECT r.id FROM agent_runtime_locks r WHERE r.agent_id=$1 AND r.id=$2 FOR UPDATE SKIP LOCKED +` + +type TryLockExecutionRuntimeParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) TryLockExecutionRuntime(ctx context.Context, db DBTX, arg TryLockExecutionRuntimeParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, tryLockExecutionRuntime, arg.AgentID, arg.ID) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_tools.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_tools.sql.go new file mode 100644 index 000000000..a4ffd4b0b --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_tools.sql.go @@ -0,0 +1,222 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: commands_tools.sql + +package executiondb + +import ( + "context" + "encoding/json" + "time" + + "github.com/google/uuid" +) + +const closeExecutionToolInteractions = `-- name: CloseExecutionToolInteractions :many +UPDATE agent_interactions SET state='canceled',resolution=jsonb_build_object('reason','tool_call_completed'),resolved_at=statement_timestamp() +WHERE agent_id=$1 AND tool_call_id=$2 AND state='open' +RETURNING id,interaction_kind,state +` + +type CloseExecutionToolInteractionsParams struct { + AgentID uuid.UUID + ToolID uuid.UUID +} + +type CloseExecutionToolInteractionsRow struct { + ID uuid.UUID + InteractionKind string + State string +} + +func (q *Queries) CloseExecutionToolInteractions(ctx context.Context, db DBTX, arg CloseExecutionToolInteractionsParams) ([]CloseExecutionToolInteractionsRow, error) { + rows, err := db.Query(ctx, closeExecutionToolInteractions, arg.AgentID, arg.ToolID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []CloseExecutionToolInteractionsRow{} + for rows.Next() { + var i CloseExecutionToolInteractionsRow + if err := rows.Scan(&i.ID, &i.InteractionKind, &i.State); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const executionToolOwner = `-- name: ExecutionToolOwner :one +SELECT EXISTS(SELECT 1 FROM agent_interactions i WHERE i.agent_id=$1 AND i.tool_call_id=$2 + AND i.id=$3 AND i.interaction_kind='question' AND i.state='open') AS question, + EXISTS(SELECT 1 FROM processes p WHERE p.agent_id=$1 AND p.tool_call_id=$2 + AND p.id=$4 AND p.state IN ('queued','starting','running')) AS process, + EXISTS(SELECT 1 FROM process_actions p WHERE p.agent_id=$1 AND p.tool_call_id=$2 + AND p.id=$5 AND p.state IN ('queued','accepted')) AS action +` + +type ExecutionToolOwnerParams struct { + AgentID uuid.UUID + ToolID uuid.UUID + InteractionID *uuid.UUID + ProcessID *uuid.UUID + ActionID *uuid.UUID +} + +type ExecutionToolOwnerRow struct { + Question bool + Process bool + Action bool +} + +func (q *Queries) ExecutionToolOwner(ctx context.Context, db DBTX, arg ExecutionToolOwnerParams) (ExecutionToolOwnerRow, error) { + row := db.QueryRow(ctx, executionToolOwner, + arg.AgentID, + arg.ToolID, + arg.InteractionID, + arg.ProcessID, + arg.ActionID, + ) + var i ExecutionToolOwnerRow + err := row.Scan(&i.Question, &i.Process, &i.Action) + return i, err +} + +const insertExecutionToolResult = `-- name: InsertExecutionToolResult :one +INSERT INTO tool_call_results(agent_id,tool_call_id,outcome,completed_at) +SELECT c.agent_id,c.id,$1,statement_timestamp() FROM tool_calls c +WHERE c.agent_id=$2 AND c.id=$3 AND c.state='completed' +RETURNING id +` + +type InsertExecutionToolResultParams struct { + Outcome string + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) InsertExecutionToolResult(ctx context.Context, db DBTX, arg InsertExecutionToolResultParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, insertExecutionToolResult, arg.Outcome, arg.AgentID, arg.ID) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const readExecutionTool = `-- name: ReadExecutionTool :one +SELECT c.id,c.agent_id,c.model_output_id,c.provider_call_id,c.name,c.input,c.type,c.state,c.runtime_lock_id, + m.turn_id,m.id AS context_id,e.id AS source_event_id,e.sequence AS source_sequence, + r.id AS result_id,r.outcome,r.completed_at,re.id AS result_event_id,re.sequence AS result_sequence +FROM tool_calls c JOIN model_outputs o ON o.agent_id=c.agent_id AND o.id=c.model_output_id +JOIN model_call_contexts m ON m.agent_id=o.agent_id AND m.id=o.model_call_context_id +JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id +LEFT JOIN agent_events re ON re.agent_id=r.agent_id AND re.tool_call_result_id=r.id +WHERE c.agent_id=$1 AND c.id=$2 +` + +type ReadExecutionToolParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +type ReadExecutionToolRow struct { + ID uuid.UUID + AgentID uuid.UUID + ModelOutputID uuid.UUID + ProviderCallID string + Name string + Input json.RawMessage + Type string + State string + RuntimeLockID *uuid.UUID + TurnID uuid.UUID + ContextID uuid.UUID + SourceEventID uuid.UUID + SourceSequence int64 + ResultID *uuid.UUID + Outcome *string + CompletedAt *time.Time + ResultEventID *uuid.UUID + ResultSequence *int64 +} + +func (q *Queries) ReadExecutionTool(ctx context.Context, db DBTX, arg ReadExecutionToolParams) (ReadExecutionToolRow, error) { + row := db.QueryRow(ctx, readExecutionTool, arg.AgentID, arg.ID) + var i ReadExecutionToolRow + err := row.Scan( + &i.ID, + &i.AgentID, + &i.ModelOutputID, + &i.ProviderCallID, + &i.Name, + &i.Input, + &i.Type, + &i.State, + &i.RuntimeLockID, + &i.TurnID, + &i.ContextID, + &i.SourceEventID, + &i.SourceSequence, + &i.ResultID, + &i.Outcome, + &i.CompletedAt, + &i.ResultEventID, + &i.ResultSequence, + ) + return i, err +} + +const readExecutionToolDispatch = `-- name: ReadExecutionToolDispatch :one +SELECT state FROM tool_calls WHERE agent_id=$1 AND id=$2 +` + +type ReadExecutionToolDispatchParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) ReadExecutionToolDispatch(ctx context.Context, db DBTX, arg ReadExecutionToolDispatchParams) (string, error) { + row := db.QueryRow(ctx, readExecutionToolDispatch, arg.AgentID, arg.ID) + var state string + err := row.Scan(&state) + return state, err +} + +const transitionExecutionTool = `-- name: TransitionExecutionTool :execrows +UPDATE tool_calls c SET state=$1,runtime_lock_id=$2 +WHERE c.agent_id=$3 AND c.id=$4 AND c.state=$5 + AND c.runtime_lock_id IS NOT DISTINCT FROM $6::uuid + AND ($7::uuid IS NULL OR EXISTS(SELECT 1 FROM agent_runtime_locks r + WHERE r.agent_id=c.agent_id AND r.id=$7 AND r.cancel_requested_at IS NULL + AND r.lease_expires_at>statement_timestamp())) +` + +type TransitionExecutionToolParams struct { + NextState string + NextRuntimeID *uuid.UUID + AgentID uuid.UUID + ID uuid.UUID + PreviousState string + PreviousRuntimeID *uuid.UUID + FenceID *uuid.UUID +} + +func (q *Queries) TransitionExecutionTool(ctx context.Context, db DBTX, arg TransitionExecutionToolParams) (int64, error) { + result, err := db.Exec(ctx, transitionExecutionTool, + arg.NextState, + arg.NextRuntimeID, + arg.AgentID, + arg.ID, + arg.PreviousState, + arg.PreviousRuntimeID, + arg.FenceID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/contexts.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/contexts.sql.go new file mode 100644 index 000000000..dab4f5c5f --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/contexts.sql.go @@ -0,0 +1,114 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: contexts.sql + +package executiondb + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +const reconstructExecutionContexts = `-- name: ReconstructExecutionContexts :many +SELECT c.id,c.agent_id,c.turn_id,c.operation_kind,c.attempt_number,c.input_event_sequence, + c.source_event_sequence_end,c.state,c.recovery_kind,c.retry_at,c.opening_input_ids,c.opening_event_sequence +FROM model_call_contexts c WHERE c.agent_id=$1 +ORDER BY c.input_event_sequence,c.attempt_number +` + +type ReconstructExecutionContextsParams struct { + AgentID uuid.UUID +} + +type ReconstructExecutionContextsRow struct { + ID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + OperationKind string + AttemptNumber int32 + InputEventSequence int64 + SourceEventSequenceEnd *int64 + State string + RecoveryKind *string + RetryAt *time.Time + OpeningInputIds []uuid.UUID + OpeningEventSequence int64 +} + +func (q *Queries) ReconstructExecutionContexts(ctx context.Context, db DBTX, arg ReconstructExecutionContextsParams) ([]ReconstructExecutionContextsRow, error) { + rows, err := db.Query(ctx, reconstructExecutionContexts, arg.AgentID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ReconstructExecutionContextsRow{} + for rows.Next() { + var i ReconstructExecutionContextsRow + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.TurnID, + &i.OperationKind, + &i.AttemptNumber, + &i.InputEventSequence, + &i.SourceEventSequenceEnd, + &i.State, + &i.RecoveryKind, + &i.RetryAt, + &i.OpeningInputIds, + &i.OpeningEventSequence, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const validateExecutionLineage = `-- name: ValidateExecutionLineage :one +WITH snapshots AS ( + SELECT c.opening_input_ids,c.opening_event_sequence,c.input_event_sequence AS watermark + FROM model_call_contexts c WHERE c.agent_id=$1 + UNION ALL + SELECT i.opening_input_ids,i.opening_event_sequence,e.sequence + FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id + WHERE i.agent_id=$1 AND i.input_kind='config_change' AND i.state='resolved' + UNION ALL + SELECT c.opening_input_ids,c.opening_event_sequence,e.sequence + FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id + WHERE c.agent_id=$1 +) +SELECT coalesce(NOT EXISTS ( + SELECT 1 FROM model_call_contexts c WHERE c.agent_id=$1 + AND c.turn_id IS DISTINCT FROM (SELECT e.turn_id FROM agent_events e + WHERE e.agent_id=c.agent_id AND e.is_opening_event AND e.sequence<=c.input_event_sequence + ORDER BY e.sequence DESC LIMIT 1) +) AND NOT EXISTS ( + SELECT 1 FROM snapshots s WHERE EXISTS ( + SELECT 1 FROM ( + SELECT o.ordinality,e.sequence,lag(e.sequence) OVER (ORDER BY o.ordinality) AS previous + FROM unnest(s.opening_input_ids) WITH ORDINALITY o(id,ordinality) + LEFT JOIN agent_inputs i ON i.agent_id=$1 AND i.id=o.id AND i.state='resolved' AND i.input_kind='content' + LEFT JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id AND e.is_opening_event + ) opening WHERE sequence IS NULL OR sequence>s.watermark + OR previous>=sequence OR (ordinality=1 AND sequence<>s.opening_event_sequence) + ) +),false)::boolean AS valid +` + +type ValidateExecutionLineageParams struct { + AgentID uuid.UUID +} + +func (q *Queries) ValidateExecutionLineage(ctx context.Context, db DBTX, arg ValidateExecutionLineageParams) (bool, error) { + row := db.QueryRow(ctx, validateExecutionLineage, arg.AgentID) + var valid bool + err := row.Scan(&valid) + return valid, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/db_gen.go b/internal/storage/internal/agentexecution/internal/executiondb/db_gen.go new file mode 100644 index 000000000..cdb4e5e5d --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/db_gen.go @@ -0,0 +1,26 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 + +package executiondb + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type DBTX interface { + Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error) + Query(context.Context, string, ...interface{}) (pgx.Rows, error) + QueryRow(context.Context, string, ...interface{}) pgx.Row + SendBatch(context.Context, *pgx.Batch) pgx.BatchResults +} + +func New() *Queries { + return &Queries{} +} + +type Queries struct { +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/inputs.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/inputs.sql.go new file mode 100644 index 000000000..9132c1f03 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/inputs.sql.go @@ -0,0 +1,107 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: inputs.sql + +package executiondb + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +const captureExecutionOpening = `-- name: CaptureExecutionOpening :many +WITH boundary AS MATERIALIZED ( + SELECT greatest( + coalesce((SELECT output.sequence FROM agent_events output WHERE output.agent_id=$1 + AND output.event_kind='model_output' AND output.sequence<=$2 + ORDER BY output.sequence DESC LIMIT 1),0), + $3::bigint + )::bigint AS sequence +), unanswered AS MATERIALIZED ( + SELECT i.id,e.sequence,e.created_at FROM boundary b + JOIN agent_events e ON e.agent_id=$1 AND e.is_opening_event + AND e.sequence>b.sequence AND e.sequence<=$2 + JOIN agent_inputs i ON i.project_id=$4 AND i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id +) +SELECT id,sequence,created_at FROM unanswered +UNION ALL +SELECT i.id,e.sequence,e.created_at FROM agent_events e +JOIN agent_inputs i ON i.project_id=$4 AND i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id +WHERE e.agent_id=$1 AND e.turn_id=$5 AND e.is_opening_event + AND e.sequence<=$2 AND NOT EXISTS(SELECT 1 FROM unanswered) +ORDER BY sequence +` + +type CaptureExecutionOpeningParams struct { + AgentID uuid.UUID + Watermark int64 + StopSequence int64 + ProjectID uuid.UUID + TurnID uuid.UUID +} + +type CaptureExecutionOpeningRow struct { + ID uuid.UUID + Sequence int64 + CreatedAt time.Time +} + +func (q *Queries) CaptureExecutionOpening(ctx context.Context, db DBTX, arg CaptureExecutionOpeningParams) ([]CaptureExecutionOpeningRow, error) { + rows, err := db.Query(ctx, captureExecutionOpening, + arg.AgentID, + arg.Watermark, + arg.StopSequence, + arg.ProjectID, + arg.TurnID, + ) + if err != nil { + return nil, err + } + defer rows.Close() + items := []CaptureExecutionOpeningRow{} + for rows.Next() { + var i CaptureExecutionOpeningRow + if err := rows.Scan(&i.ID, &i.Sequence, &i.CreatedAt); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const reconstructExecutionConfig = `-- name: ReconstructExecutionConfig :one +SELECT i.id FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id +WHERE i.agent_id=$1 AND e.turn_id=$2 AND i.input_kind='config_change' AND i.state='resolved' + AND e.sequence>$3 AND e.sequence>=$4 + AND e.sequence>=$5 AND $5::bigint>0 +ORDER BY e.sequence DESC LIMIT 1 +` + +type ReconstructExecutionConfigParams struct { + AgentID uuid.UUID + TurnID uuid.UUID + NormalWatermark int64 + StopSequence int64 + FirstContentSequence int64 +} + +func (q *Queries) ReconstructExecutionConfig(ctx context.Context, db DBTX, arg ReconstructExecutionConfigParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, reconstructExecutionConfig, + arg.AgentID, + arg.TurnID, + arg.NormalWatermark, + arg.StopSequence, + arg.FirstContentSequence, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go new file mode 100644 index 000000000..ca4e5cb82 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go @@ -0,0 +1,81 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: load.sql + +package executiondb + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +const loadExecutionBase = `-- name: LoadExecutionBase :one +SELECT coalesce(h.agent_id,a.id)::uuid AS agent_id,h.current_turn_id, + coalesce(h.stop_sequence,0)::bigint AS stop_sequence, + coalesce(h.answered_through_sequence,0)::bigint AS answered_through_sequence, + coalesce(h.max_normal_input_sequence,0)::bigint AS max_normal_input_sequence, + coalesce(h.max_context_input_sequence,0)::bigint AS max_context_input_sequence, + h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id, + h.pending_config_input_id,h.pending_checkpoint_id,coalesce(h.turn_continuable,false)::boolean AS turn_continuable, + coalesce(h.incomplete_tools,false)::boolean AS incomplete_tools,h.logical_ready_at, + (h.agent_id IS NOT NULL)::boolean AS head_exists,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now +FROM agents a LEFT JOIN agent_execution_state h ON h.agent_id=a.id +WHERE a.project_id=$1 AND a.id=$2 +` + +type LoadExecutionBaseParams struct { + ProjectID uuid.UUID + AgentID uuid.UUID +} + +type LoadExecutionBaseRow struct { + AgentID uuid.UUID + CurrentTurnID *uuid.UUID + StopSequence int64 + AnsweredThroughSequence int64 + MaxNormalInputSequence int64 + MaxContextInputSequence int64 + NormalContextID *uuid.UUID + CompactionContextID *uuid.UUID + PendingToolOutputID *uuid.UUID + PendingOutputLimitID *uuid.UUID + PendingConfigInputID *uuid.UUID + PendingCheckpointID *uuid.UUID + TurnContinuable bool + IncompleteTools bool + LogicalReadyAt *time.Time + HeadExists bool + RootAgentID uuid.UUID + State string + DatabaseNow time.Time +} + +func (q *Queries) LoadExecutionBase(ctx context.Context, db DBTX, arg LoadExecutionBaseParams) (LoadExecutionBaseRow, error) { + row := db.QueryRow(ctx, loadExecutionBase, arg.ProjectID, arg.AgentID) + var i LoadExecutionBaseRow + err := row.Scan( + &i.AgentID, + &i.CurrentTurnID, + &i.StopSequence, + &i.AnsweredThroughSequence, + &i.MaxNormalInputSequence, + &i.MaxContextInputSequence, + &i.NormalContextID, + &i.CompactionContextID, + &i.PendingToolOutputID, + &i.PendingOutputLimitID, + &i.PendingConfigInputID, + &i.PendingCheckpointID, + &i.TurnContinuable, + &i.IncompleteTools, + &i.LogicalReadyAt, + &i.HeadExists, + &i.RootAgentID, + &i.State, + &i.DatabaseNow, + ) + return i, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/models_gen.go b/internal/storage/internal/agentexecution/internal/executiondb/models_gen.go new file mode 100644 index 000000000..3310dd1c9 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/models_gen.go @@ -0,0 +1,5 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 + +package executiondb diff --git a/internal/storage/internal/agentexecution/internal/executiondb/outputs.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/outputs.sql.go new file mode 100644 index 000000000..3ecb8ff37 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/outputs.sql.go @@ -0,0 +1,62 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: outputs.sql + +package executiondb + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +const reconstructExecutionOutputs = `-- name: ReconstructExecutionOutputs :many +SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at +FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +WHERE o.agent_id=$1 AND e.turn_id=$2 AND e.sequence>=$3 +ORDER BY e.sequence +` + +type ReconstructExecutionOutputsParams struct { + AgentID uuid.UUID + TurnID uuid.UUID + StopSequence int64 +} + +type ReconstructExecutionOutputsRow struct { + ID uuid.UUID + ModelCallContextID uuid.UUID + StopReason string + EventID uuid.UUID + Sequence int64 + CreatedAt time.Time +} + +func (q *Queries) ReconstructExecutionOutputs(ctx context.Context, db DBTX, arg ReconstructExecutionOutputsParams) ([]ReconstructExecutionOutputsRow, error) { + rows, err := db.Query(ctx, reconstructExecutionOutputs, arg.AgentID, arg.TurnID, arg.StopSequence) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ReconstructExecutionOutputsRow{} + for rows.Next() { + var i ReconstructExecutionOutputsRow + if err := rows.Scan( + &i.ID, + &i.ModelCallContextID, + &i.StopReason, + &i.EventID, + &i.Sequence, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/publication.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/publication.sql.go new file mode 100644 index 000000000..5079fe3e8 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/publication.sql.go @@ -0,0 +1,6 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: publication.sql + +package executiondb diff --git a/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go new file mode 100644 index 000000000..d17182f11 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go @@ -0,0 +1,58 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: state.sql + +package executiondb + +import ( + "context" + + "github.com/google/uuid" +) + +const loadExecutionBoundaries = `-- name: LoadExecutionBoundaries :one +SELECT coalesce((SELECT max(e.sequence) FROM agent_events e + JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + WHERE i.project_id=$1 AND i.agent_id=$2 AND i.input_kind='control' + AND i.state='resolved' AND i.control_type='cancel_current'),0)::bigint AS stop_sequence, + coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$2 AND e.event_kind='model_output' + ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS output_sequence, + coalesce((SELECT t.id FROM agent_turns t WHERE t.agent_id=$2 ORDER BY t.turn_sequence DESC LIMIT 1), + '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS turn_id +` + +type LoadExecutionBoundariesParams struct { + ProjectID uuid.UUID + AgentID uuid.UUID +} + +type LoadExecutionBoundariesRow struct { + StopSequence int64 + OutputSequence int64 + TurnID uuid.UUID +} + +func (q *Queries) LoadExecutionBoundaries(ctx context.Context, db DBTX, arg LoadExecutionBoundariesParams) (LoadExecutionBoundariesRow, error) { + row := db.QueryRow(ctx, loadExecutionBoundaries, arg.ProjectID, arg.AgentID) + var i LoadExecutionBoundariesRow + err := row.Scan(&i.StopSequence, &i.OutputSequence, &i.TurnID) + return i, err +} + +const loadExecutionScope = `-- name: LoadExecutionScope :one +SELECT a.root_agent_id +FROM agents a WHERE a.project_id=$1 AND a.id=$2 +` + +type LoadExecutionScopeParams struct { + ProjectID uuid.UUID + ID uuid.UUID +} + +func (q *Queries) LoadExecutionScope(ctx context.Context, db DBTX, arg LoadExecutionScopeParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, loadExecutionScope, arg.ProjectID, arg.ID) + var root_agent_id uuid.UUID + err := row.Scan(&root_agent_id) + return root_agent_id, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/tools.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/tools.sql.go new file mode 100644 index 000000000..e4259c0c7 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/tools.sql.go @@ -0,0 +1,73 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: tools.sql + +package executiondb + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +const loadExecutionBatch = `-- name: LoadExecutionBatch :one +SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2) AS has_tools, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2 AND c.state<>'completed') AS incomplete, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2 AND + (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable +` + +type LoadExecutionBatchParams struct { + AgentID uuid.UUID + ModelOutputID uuid.UUID +} + +type LoadExecutionBatchRow struct { + HasTools bool + Incomplete bool + Runnable bool +} + +func (q *Queries) LoadExecutionBatch(ctx context.Context, db DBTX, arg LoadExecutionBatchParams) (LoadExecutionBatchRow, error) { + row := db.QueryRow(ctx, loadExecutionBatch, arg.AgentID, arg.ModelOutputID) + var i LoadExecutionBatchRow + err := row.Scan(&i.HasTools, &i.Incomplete, &i.Runnable) + return i, err +} + +const loadExecutionBatchCompletion = `-- name: LoadExecutionBatchCompletion :one +SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, + coalesce(max(e.sequence),0)::bigint AS last_result_sequence, + coalesce(max(r.completed_at),'epoch'::timestamptz)::timestamptz AS ready_at +FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id +LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id +WHERE c.agent_id=$1 AND c.model_output_id=$2 +` + +type LoadExecutionBatchCompletionParams struct { + AgentID uuid.UUID + ModelOutputID uuid.UUID +} + +type LoadExecutionBatchCompletionRow struct { + Calls int64 + Results int64 + Events int64 + LastResultSequence int64 + ReadyAt time.Time +} + +func (q *Queries) LoadExecutionBatchCompletion(ctx context.Context, db DBTX, arg LoadExecutionBatchCompletionParams) (LoadExecutionBatchCompletionRow, error) { + row := db.QueryRow(ctx, loadExecutionBatchCompletion, arg.AgentID, arg.ModelOutputID) + var i LoadExecutionBatchCompletionRow + err := row.Scan( + &i.Calls, + &i.Results, + &i.Events, + &i.LastResultSequence, + &i.ReadyAt, + ) + return i, err +} diff --git a/internal/storage/internal/agentexecution/internal/executiondb/turns_events.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/turns_events.sql.go new file mode 100644 index 000000000..3a7acc332 --- /dev/null +++ b/internal/storage/internal/agentexecution/internal/executiondb/turns_events.sql.go @@ -0,0 +1,79 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: turns_events.sql + +package executiondb + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +const loadExecutionTurn = `-- name: LoadExecutionTurn :one +SELECT t.id,semantic.sequence AS semantic_sequence,semantic.created_at AS semantic_time, + coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=t.agent_id AND e.turn_id=t.id AND e.is_opening_event + ORDER BY e.sequence LIMIT 1),0)::bigint AS first_opening_sequence, + coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=t.agent_id AND e.turn_id=t.id AND e.is_opening_event + ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS last_opening_sequence, + coalesce((SELECT e.sequence FROM agent_events e JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + WHERE e.agent_id=t.agent_id AND e.turn_id=t.id AND e.is_opening_event AND i.input_kind='content' + ORDER BY e.sequence LIMIT 1),0)::bigint AS first_content_sequence +FROM agent_turns t JOIN agent_events semantic ON semantic.agent_id=t.agent_id AND semantic.id=t.latest_semantic_event_id +WHERE t.agent_id=$1 AND t.id=$2 +` + +type LoadExecutionTurnParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +type LoadExecutionTurnRow struct { + ID uuid.UUID + SemanticSequence int64 + SemanticTime time.Time + FirstOpeningSequence int64 + LastOpeningSequence int64 + FirstContentSequence int64 +} + +func (q *Queries) LoadExecutionTurn(ctx context.Context, db DBTX, arg LoadExecutionTurnParams) (LoadExecutionTurnRow, error) { + row := db.QueryRow(ctx, loadExecutionTurn, arg.AgentID, arg.ID) + var i LoadExecutionTurnRow + err := row.Scan( + &i.ID, + &i.SemanticSequence, + &i.SemanticTime, + &i.FirstOpeningSequence, + &i.LastOpeningSequence, + &i.FirstContentSequence, + ) + return i, err +} + +const reconstructExecutionCheckpoint = `-- name: ReconstructExecutionCheckpoint :one +SELECT c.id FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=$1 AND e.turn_id=$2 AND e.sequence>$3 AND e.sequence>=$4 +ORDER BY e.sequence DESC LIMIT 1 +` + +type ReconstructExecutionCheckpointParams struct { + AgentID uuid.UUID + TurnID uuid.UUID + NormalWatermark int64 + StopSequence int64 +} + +func (q *Queries) ReconstructExecutionCheckpoint(ctx context.Context, db DBTX, arg ReconstructExecutionCheckpointParams) (uuid.UUID, error) { + row := db.QueryRow(ctx, reconstructExecutionCheckpoint, + arg.AgentID, + arg.TurnID, + arg.NormalWatermark, + arg.StopSequence, + ) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} diff --git a/internal/storage/internal/agentexecution/lineage_integration_test.go b/internal/storage/internal/agentexecution/lineage_integration_test.go new file mode 100644 index 000000000..e2a0a22da --- /dev/null +++ b/internal/storage/internal/agentexecution/lineage_integration_test.go @@ -0,0 +1,133 @@ +//go:build integration + +package agentexecution_test + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func TestExecutionGoverningCompactionRetainsTerminalIdentity(t *testing.T) { + f := releasedExecutionFixture(t) + agent := f.agent(t) + f.opening(t, agent, 1, uuid.Nil, "content") + f.opening(t, agent, 2, uuid.Nil, "content") + normal := f.modelContext(t, agent, 2, "normal", nil, 1) + f.fail(t, normal, "compact") + larger := f.modelContext(t, agent, 2, "compaction", new(int64(2)), 1) + f.fail(t, larger, "reduce_compaction_source") + smaller := f.modelContext(t, agent, 2, "compaction", new(int64(1)), 1) + f.fail(t, smaller, "retry") + latest := f.modelContext(t, agent, 2, "compaction", new(int64(1)), 2) + _, err := f.pool.Exec( + t.Context(), + `UPDATE model_call_contexts SET state='canceled',error_kind='cancel',completed_at=$2 WHERE id=$1`, + latest, + f.now, + ) + require.NoError(t, err) + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + _, h := f.handle(t, agent) + snapshot, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + rebuilt, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + require.Equal(t, snapshot.Head, rebuilt.Head) + require.Equal(t, latest, rebuilt.Head.CompactionContextID) + require.Nil(t, rebuilt.Selection.Model) + require.Equal(t, agentexecution.WaitIdle, rebuilt.Selection.Wait) +} + +func TestExecutionEmptyAndUnansweredConfigLineage(t *testing.T) { + f := releasedExecutionFixture(t) + empty, unanswered := f.agent(t), f.agent(t) + f.opening(t, empty, 1, uuid.Nil, "config_change") + _, content := f.opening(t, unanswered, 1, uuid.Nil, "content") + f.opening(t, unanswered, 2, uuid.Nil, "config_change") + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + for _, id := range []uuid.UUID{empty, unanswered} { + u, h := f.handle(t, id) + snapshot, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + if id == empty { + require.Nil(t, snapshot.Selection.Model) + } else { + require.Equal(t, agentexecution.OriginInitial, snapshot.Selection.Model.Origin) + require.Equal(t, []uuid.UUID{content}, snapshot.Selection.Model.Opening.InputIDs) + } + + require.NoError(t, u.Rollback(t.Context())) + } +} + +func TestExecutionCutoverRejectsOrphanCompaction(t *testing.T) { + f := releasedExecutionFixture(t) + agent := f.agent(t) + f.opening(t, agent, 1, uuid.Nil, "content") + f.modelContext(t, agent, 1, "compaction", new(int64(1)), 1) + require.ErrorContains(t, applyExecutionCutover(t.Context(), f.pool), "unrepresented governing context") +} + +func TestExecutionCutoverReconcilesOwnershipAgeAndDeadlines(t *testing.T) { + f := releasedExecutionFixture(t) + ready, owned, retry, archived := f.agent(t), f.agent(t), f.agent(t), f.agent(t) + for _, agent := range []uuid.UUID{ready, owned, retry, archived} { + f.opening(t, agent, 1, uuid.Nil, "content") + _, err := f.pool.Exec(t.Context(), `INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) + VALUES($1,$2::timestamptz-interval '1 hour',$2)`, agent, f.now) + require.NoError(t, err) + } + model := f.modelContext(t, retry, 1, "normal", nil, 1) + f.fail(t, model, "retry") + _, err := f.pool.Exec( + t.Context(), + `INSERT INTO agent_runtime_locks(agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + VALUES($1,$2,$3,$3,$3::timestamptz+interval '90 seconds')`, + owned, + uuid.New(), + f.now, + ) + require.NoError(t, err) + _, err = f.pool.Exec( + t.Context(), + `UPDATE agents SET state='archived',archived_at=$2 WHERE id=$1`, + archived, + f.now, + ) + require.NoError(t, err) + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + var queueCount int + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_wakeups`).Scan(&queueCount)) + require.Equal(t, 2, queueCount) + var age, deadline time.Time + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, ready). + Scan(&age), + ) + require.Equal(t, f.now.Add(-time.Hour), age.UTC()) + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT ready_at FROM agent_wakeups WHERE agent_id=$1`, retry). + Scan(&deadline), + ) + require.Equal(t, f.now.Add(time.Hour), deadline.UTC()) + var ownedReady bool + require.NoError( + t, + f.pool.QueryRow(t.Context(), ` +SELECT logical_ready_at IS NOT NULL FROM agent_execution_state WHERE agent_id=$1 +`, owned). + Scan(&ownedReady), + ) + require.True(t, ownedReady) + _, h := f.handle(t, archived) + snapshot, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + require.Equal(t, agentexecution.WaitArchived, snapshot.Selection.Wait) + require.Nil(t, snapshot.Head.LogicalReadyAt) +} diff --git a/internal/storage/internal/agentexecution/linearity_integration_test.go b/internal/storage/internal/agentexecution/linearity_integration_test.go new file mode 100644 index 000000000..a3a817435 --- /dev/null +++ b/internal/storage/internal/agentexecution/linearity_integration_test.go @@ -0,0 +1,133 @@ +//go:build integration + +package agentexecution_test + +import ( + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/stretchr/testify/require" +) + +func TestExecutionBulkAndLoopLinearity(t *testing.T) { + f := migratedExecutionFixture(t) + ctx := t.Context() + for _, mode := range []string{"bulk", "loop"} { + t.Run(mode, func(t *testing.T) { + var baseline time.Duration + for _, n := range []int{1000, 10000} { + t.Run(fmt.Sprint(n), func(t *testing.T) { + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(ctx) + require.NoError(t, err) + defer func() { _ = u.Rollback(ctx) }() + started := time.Now() + ids := make([]uuid.UUID, n) + for i := range ids { + ids[i] = uuid.New() + } + if mode == "bulk" { + _, err = u.DB(). + Exec(ctx, ` +INSERT INTO agents(id,root_agent_id,org_id,project_id,state,current_config_id,created_at,updated_at) + SELECT id,id,$2,$3,'active',$4,$5,$5 FROM unnest($1::uuid[]) id`, ids, f.org, f.project, f.config, f.now) + require.NoError(t, err) + _, err = u.DB(). + Exec(ctx, `INSERT INTO agent_execution_state(agent_id,turn_continuable,incomplete_tools) + SELECT id,false,false FROM unnest($1::uuid[]) id`, ids) + require.NoError(t, err) + } else { + for _, id := range ids { + _, err = u.DB().Exec(ctx, ` +INSERT INTO agents(id,root_agent_id,org_id,project_id,state,current_config_id,created_at,updated_at) + VALUES($1,$1,$2,$3,'active',$4,$5,$5)`, id, f.org, f.project, f.config, f.now) + require.NoError(t, err) + require.NoError(t, u.CreatedAgent(f.project, id, uuid.Nil)) + _, err = u.DB().Exec(ctx, ` +INSERT INTO agent_execution_state(agent_id,turn_continuable,incomplete_tools) VALUES($1,false,false) +`, id) + require.NoError(t, err) + } + } + require.NoError(t, u.Commit(ctx, "linearity fixture")) + elapsed := time.Since(started) + t.Logf("%s %d agents: %s (%s/agent)", mode, n, elapsed, elapsed/time.Duration(n)) + var count int + require.NoError( + t, + f.pool.QueryRow(ctx, `SELECT count(*) FROM agent_execution_state WHERE agent_id=ANY($1)`, ids). + Scan(&count), + ) + require.Equal(t, n, count) + if n == 1000 { + baseline = elapsed + } else { + require.Less(t, elapsed, baseline*25+time.Second) + } + }) + } + }) + } +} + +func TestExecutionReconstructionTenThousandAgents(t *testing.T) { + f := releasedExecutionFixture(t) + ctx := t.Context() + _, err := f.pool.Exec( + ctx, + `INSERT INTO agents(org_id,project_id,state,current_config_id,created_at,updated_at) + SELECT $1,$2,'active',$3,$4,$4 FROM generate_series(1,10000)`, + f.org, + f.project, + f.config, + f.now, + ) + require.NoError(t, err) + started := time.Now() + require.NoError(t, applyExecutionCutover(ctx, f.pool)) + t.Logf("released-51 cutover backfill of 10000 agents: %s", time.Since(started)) + rows, err := f.pool.Query(ctx, `SELECT id FROM agents ORDER BY id`) + require.NoError(t, err) + var refs []lifecyclelock.AgentRef + for rows.Next() { + var id uuid.UUID + require.NoError(t, rows.Scan(&id)) + refs = append(refs, lifecyclelock.AgentRef{ProjectID: f.project, AgentID: id}) + } + require.NoError(t, rows.Err()) + rows.Close() + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(ctx) + require.NoError(t, err) + defer func() { _ = u.Rollback(ctx) }() + started = time.Now() + plan, err := u.PlanAgents(ctx, refs, agentexecution.RepairAuthority{}) + require.NoError(t, err) + require.NoError(t, u.LockAgents(ctx, plan)) + var firstThousand time.Duration + for i, ref := range refs { + h, handleErr := u.Agent( + agentexecution.AgentRoute{ + CellID: cell.ID(), + ProjectID: f.project, + AgentID: ref.AgentID, + RootAgentID: ref.AgentID, + }, + ) + require.NoError(t, handleErr) + snapshot, loadErr := h.ReconstructExecution(ctx) + require.NoError(t, loadErr) + require.Equal(t, agentexecution.WaitIdle, snapshot.Selection.Wait) + if i == 999 { + firstThousand = time.Since(started) + started = time.Now() + } + } + remaining := time.Since(started) + t.Logf("lock plan and first 1000 reconstructions: %s; remaining 9000: %s", firstThousand, remaining) + require.Less(t, remaining, firstThousand*20+time.Second) +} diff --git a/internal/storage/internal/agentexecution/load.go b/internal/storage/internal/agentexecution/load.go new file mode 100644 index 000000000..544db38dd --- /dev/null +++ b/internal/storage/internal/agentexecution/load.go @@ -0,0 +1,190 @@ +package agentexecution + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" +) + +func nullableID(id uuid.UUID) *uuid.UUID { + if id == uuid.Nil { + return nil + } + return &id +} + +func loadExecution( + ctx context.Context, + db executiondb.DBTX, + route AgentRoute, + override *ExecutionHead, +) (ExecutionSnapshot, error) { + q := executiondb.New() + if override != nil { + return loadExecutionFacts(ctx, db, route, *override) + } + base, err := q.LoadExecutionBase(ctx, db, executiondb.LoadExecutionBaseParams{ + ProjectID: route.ProjectID, AgentID: route.AgentID, + }) + if err != nil { + return ExecutionSnapshot{}, fmt.Errorf("load execution base: %w", err) + } + if !base.HeadExists { + return ExecutionSnapshot{}, fmt.Errorf("load execution head: %w", pgx.ErrNoRows) + } + return loadExecutionFacts(ctx, db, route, executionHead(base)) +} + +func executionFactsParams(route AgentRoute, head ExecutionHead) executiondb.LoadExecutionFactsParams { + return executiondb.LoadExecutionFactsParams{AgentID: route.AgentID, ProjectID: route.ProjectID, + TurnID: nullableID( + head.CurrentTurnID, + ), StopSequence: head.StopSequence, MaxContextInputSequence: head.MaxContextInputSequence, + NormalContextID: nullableID( + head.NormalContextID, + ), CompactionContextID: nullableID(head.CompactionContextID), + ToolOutputID: nullableID( + head.PendingToolOutputID, + ), LimitOutputID: nullableID(head.PendingOutputLimitID), + ConfigID: nullableID(head.PendingConfigInputID), CheckpointID: nullableID(head.PendingCheckpointID)} +} + +func loadExecutionFacts( + ctx context.Context, + db executiondb.DBTX, + route AgentRoute, + head ExecutionHead, +) (ExecutionSnapshot, error) { + batch := executiondb.New(). + LoadExecutionFacts(ctx, db, []executiondb.LoadExecutionFactsParams{executionFactsParams(route, head)}) + var snapshot ExecutionSnapshot + var err error + batch.Query(func(_ int, rows []executiondb.LoadExecutionFactsRow, queryErr error) { + if queryErr != nil { + err = queryErr + return + } + snapshot, err = executionSnapshot(route, head, rows) + }) + return snapshot, errors.Join(err, batch.Close()) +} + +func executionSnapshot( + route AgentRoute, + head ExecutionHead, + rows []executiondb.LoadExecutionFactsRow, +) (ExecutionSnapshot, error) { + view := ExecutionView{AgentID: route.AgentID, StopSequence: head.StopSequence, + MaxNormalInputSequence: head.MaxNormalInputSequence, MaxContextInputSequence: head.MaxContextInputSequence} + var now time.Time + for _, row := range rows { + if row.Kind != "scope" { + continue + } + if row.ID != route.AgentID || row.RootAgentID != route.RootAgentID || + row.TurnID != head.CurrentTurnID { + return ExecutionSnapshot{}, fmt.Errorf("%w: execution scope", ErrInvalidState) + } + view.State = AgentState(row.State) + now = row.DatabaseNow + view.Inputs = InputAvailability{Steering: row.Steering, Queued: row.Queued} + if row.TurnID != uuid.Nil { + view.Turn = &Turn{ID: row.TurnID, FirstOpeningSequence: row.FirstOpeningSequence, + LastOpeningSequence: row.LastOpeningSequence, FirstContentSequence: row.FirstContentSequence, + LatestSemantic: EventBoundary{Sequence: row.Sequence, Time: row.EventTime}} + } + } + if now.IsZero() { + return ExecutionSnapshot{}, fmt.Errorf("%w: missing agent", ErrInvalidState) + } + contexts := make(map[uuid.UUID]*ModelContext) + var batch *ToolBatch + for _, row := range rows { + switch row.Kind { + case "context": + contexts[row.ID] = &ModelContext{ID: row.ID, AgentID: route.AgentID, TurnID: row.TurnID, + Operation: Operation(row.Operation), Attempt: row.Attempt, InputEventSequence: row.Watermark, + SourceEventSequenceEnd: row.SourceEnd, State: ContextState(row.State), Recovery: RecoveryKind(row.Recovery), + RetryAt: row.RetryAt, Opening: Opening{InputIDs: row.OpeningIds, EventSequence: row.OpeningSequence}} + case "batch": + if !row.HasTools { + return ExecutionSnapshot{}, fmt.Errorf("%w: pending output has no tools", ErrInvalidState) + } + batch = &ToolBatch{HasIncomplete: row.Incomplete, HasRunnable: row.Runnable} + if !row.Incomplete { + if row.Calls == 0 || row.Calls != row.Results || row.Calls != row.Events { + return ExecutionSnapshot{}, fmt.Errorf("%w: incomplete result authority", ErrInvalidState) + } + batch.Completion = &BatchCompletion{LastResultSequence: row.Sequence, ReadyAt: row.EventTime} + } + case "config", "checkpoint": + work := &BoundaryWork{ID: row.ID, AgentID: route.AgentID, TurnID: row.TurnID, + Event: EventBoundary{Sequence: row.Sequence, Time: row.EventTime}, + Opening: Opening{InputIDs: row.OpeningIds, EventSequence: row.OpeningSequence}} + if row.Kind == "config" { + view.Config = work + } else { + view.Checkpoint = work + } + case "opening": + if view.Turn == nil || head.MaxContextInputSequence >= view.Turn.FirstOpeningSequence || + head.StopSequence > view.Turn.FirstOpeningSequence { + return ExecutionSnapshot{}, ErrInvalidState + } + if len(view.Turn.InitialOpening.InputIDs) == 0 { + view.Turn.InitialOpening.EventSequence = row.Sequence + view.Turn.InitialReadyAt = row.EventTime + } + view.Turn.InitialOpening.InputIDs = append(view.Turn.InitialOpening.InputIDs, row.ID) + if row.EventTime.Before(view.Turn.InitialReadyAt) { + view.Turn.InitialReadyAt = row.EventTime + } + } + } + view.NormalContext = contexts[head.NormalContextID] + view.CompactionContext = contexts[head.CompactionContextID] + for _, row := range rows { + if row.Kind != "output" { + continue + } + source := contexts[row.ContextID] + if source == nil { + return ExecutionSnapshot{}, fmt.Errorf("%w: missing output context", ErrInvalidState) + } + output := ModelOutput{ + ID: row.ID, + EventID: row.EventID, + Event: EventBoundary{Sequence: row.Sequence, Time: row.EventTime}, + Context: *source, + } + if row.ID == head.PendingToolOutputID && batch != nil { + batch.Output = output + view.ToolBatch = batch + } else if row.ID == head.PendingOutputLimitID && row.StopReason == "max_tokens" { + view.OutputLimit = &output + } else { + return ExecutionSnapshot{}, fmt.Errorf("%w: pending output", ErrInvalidState) + } + } + if head.NormalContextID != uuid.Nil && view.NormalContext == nil || + head.CompactionContextID != uuid.Nil && view.CompactionContext == nil || + head.PendingToolOutputID != uuid.Nil && view.ToolBatch == nil || + head.PendingOutputLimitID != uuid.Nil && view.OutputLimit == nil || + head.PendingConfigInputID != uuid.Nil && view.Config == nil || + head.PendingCheckpointID != uuid.Nil && view.Checkpoint == nil { + return ExecutionSnapshot{}, fmt.Errorf("%w: missing execution pointer", ErrInvalidState) + } + selected, err := SelectNext(view, now) + if err != nil { + return ExecutionSnapshot{}, err + } + head.TurnContinuable = selected.TurnContinuable + head.IncompleteTools = selected.IncompleteTools + head.LogicalReadyAt = selected.LogicalReadyAt + return ExecutionSnapshot{Head: head, View: view, Selection: selected}, nil +} diff --git a/internal/storage/internal/agentexecution/load_count_integration_test.go b/internal/storage/internal/agentexecution/load_count_integration_test.go new file mode 100644 index 000000000..18835602e --- /dev/null +++ b/internal/storage/internal/agentexecution/load_count_integration_test.go @@ -0,0 +1,65 @@ +//go:build integration + +package agentexecution + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type countingLoadTx struct { + pgx.Tx + statements int + batches int + queries int +} + +func (tx *countingLoadTx) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { + tx.statements++ + tx.queries++ + return tx.Tx.Exec(ctx, sql, args...) +} + +func (tx *countingLoadTx) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { + tx.statements++ + tx.queries++ + rows, err := tx.Tx.Query(ctx, sql, args...) + if err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + rows.Close() + return nil, err + } + return rows, nil +} + +func (tx *countingLoadTx) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { + tx.statements++ + tx.queries++ + return tx.Tx.QueryRow(ctx, sql, args...) +} + +func CountOrdinaryLoad(ctx context.Context, h *Handle) (ExecutionSnapshot, int, error) { + tx := &countingLoadTx{Tx: h.unit.tx} + h.unit.tx = tx + defer func() { h.unit.tx = tx.Tx }() + snapshot, err := h.LoadExecution(ctx) + return snapshot, tx.statements, err +} + +func (tx *countingLoadTx) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { + tx.statements += batch.Len() + tx.batches++ + return tx.Tx.SendBatch(ctx, batch) +} + +func CountCommit(ctx context.Context, u *Unit) (int, int, error) { + tx := &countingLoadTx{Tx: u.tx} + u.tx = tx + defer func() { u.tx = tx.Tx }() + err := u.Commit(ctx, "count publication") + return tx.statements, tx.batches + tx.queries, err +} diff --git a/internal/storage/internal/agentexecution/load_integration_test.go b/internal/storage/internal/agentexecution/load_integration_test.go new file mode 100644 index 000000000..32a3d3570 --- /dev/null +++ b/internal/storage/internal/agentexecution/load_integration_test.go @@ -0,0 +1,48 @@ +//go:build integration + +package agentexecution_test + +import ( + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func TestOrdinaryExecutionLoadStatementBound(t *testing.T) { + f := releasedExecutionFixture(t) + ids := make(map[string]uuid.UUID) + for _, name := range []string{"idle", "initial", "incomplete", "complete", "retry"} { + a := f.agent(t) + ids[name] = a + if name == "idle" { + continue + } + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + if name == "initial" { + continue + } + c := f.modelContext(t, a, 1, "normal", nil, 1) + if name == "retry" { + f.fail(t, c, "retry") + continue + } + _, call := f.output(t, a, turn, c, 2, "tool_use", "awaiting_authorization") + if name == "complete" { + f.complete(t, a, turn, call, 3) + } + } + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + for name, id := range ids { + t.Run(name, func(t *testing.T) { + _, h := f.handle(t, id) + snapshot, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + require.NoError(t, err) + require.Equal(t, 2, count) + if snapshot.View.Turn != nil && name != "initial" { + require.Empty(t, snapshot.View.Turn.InitialOpening.InputIDs) + } + }) + } +} diff --git a/internal/storage/internal/agentexecution/locks.go b/internal/storage/internal/agentexecution/locks.go new file mode 100644 index 000000000..194d5a295 --- /dev/null +++ b/internal/storage/internal/agentexecution/locks.go @@ -0,0 +1,254 @@ +package agentexecution + +import ( + "context" + "errors" + "fmt" + "slices" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type MutationAuthority interface{ mutationAuthority() } + +type RuntimeAuthority struct { + AgentID uuid.UUID + RuntimeLockID uuid.UUID +} + +type IngressAuthority struct{} +type LifecycleAuthority struct{} +type RepairAuthority struct{} +type ExternalAuthority struct{} + +func (RuntimeAuthority) mutationAuthority() {} +func (IngressAuthority) mutationAuthority() {} +func (LifecycleAuthority) mutationAuthority() {} +func (RepairAuthority) mutationAuthority() {} +func (ExternalAuthority) mutationAuthority() {} + +type MutationPlan struct { + agents []AgentRoute + authority MutationAuthority + unit *Unit +} + +func (u *Unit) PlanAgents( + ctx context.Context, refs []lifecyclelock.AgentRef, authority MutationAuthority, +) (MutationPlan, error) { + if err := u.active(); err != nil { + return MutationPlan{}, err + } + plan := MutationPlan{authority: authority, unit: u} + q := executiondb.New() + for _, ref := range refs { + if h := u.handles[ref.AgentID]; h != nil && h.route.ProjectID == ref.ProjectID { + plan.agents = append(plan.agents, h.route) + continue + } + root, err := q.LoadExecutionScope(ctx, u.DB(), executiondb.LoadExecutionScopeParams{ + ProjectID: ref.ProjectID, ID: ref.AgentID, + }) + if errors.Is(err, pgx.ErrNoRows) { + return MutationPlan{}, storeerr.ErrNotFound + } + if err != nil { + return MutationPlan{}, err + } + plan.agents = append(plan.agents, AgentRoute{ + CellID: u.cell.id, ProjectID: ref.ProjectID, AgentID: ref.AgentID, RootAgentID: root, + }) + } + return plan, nil +} + +func (u *Unit) LockAgents(ctx context.Context, plan MutationPlan) error { + _, err := u.lockAgents(ctx, plan, false) + return err +} + +func (u *Unit) TryLockAgents(ctx context.Context, plan MutationPlan) (bool, error) { + return u.lockAgents(ctx, plan, true) +} + +func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked bool) (bool, error) { + if err := u.active(); err != nil { + return false, err + } + switch plan.authority.(type) { + case RuntimeAuthority, IngressAuthority, LifecycleAuthority, RepairAuthority, ExternalAuthority: + default: + return false, errors.New("agent mutation authority is required") + } + if plan.unit != u { + return false, fmt.Errorf("%w: plan belongs to a different unit", ErrLockPlan) + } + routes := slices.Clone(plan.agents) + slices.SortFunc(routes, compareRoutes) + routes = slices.Compact(routes) + for _, route := range routes { + if route.CellID != u.cell.id || route.ProjectID == uuid.Nil || + route.AgentID == uuid.Nil || route.RootAgentID == uuid.Nil { + return false, fmt.Errorf("%w: invalid agent route", ErrLockPlan) + } + if h := u.handles[route.AgentID]; h != nil { + if h.route != route { + return false, fmt.Errorf("%w: route changed", ErrLockPlan) + } + continue + } + if u.runtimeLocked { + return false, fmt.Errorf("%w: agent acquired after runtime", ErrLockPlan) + } + for _, held := range u.handles { + _, excluded := u.exclusiveProjects[route.ProjectID] + if held.created || excluded { + continue + } + if compareRoutes(route, held.route) < 0 { + return false, fmt.Errorf("%w: agent order", ErrLockPlan) + } + } + } + q := dbsqlc.New(u.DB()) + for _, route := range routes { + if u.handles[route.AgentID] != nil { + continue + } + var err error + if skipLocked { + _, err = q.TryLockAgentInProject(ctx, dbsqlc.TryLockAgentInProjectParams{ + ProjectID: route.ProjectID, ID: route.AgentID, + }) + } else { + _, err = q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: route.ProjectID, ID: route.AgentID, + }) + } + if errors.Is(err, pgx.ErrNoRows) { + if skipLocked { + return false, nil + } + return false, storeerr.ErrNotFound + } + if err != nil { + return false, err + } + u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} + } + if owned, ok := plan.authority.(RuntimeAuthority); ok { + h := u.handles[owned.AgentID] + if h == nil || owned.RuntimeLockID == uuid.Nil { + return false, storeerr.ErrRuntimeLockInactive + } + if _, err := q.LockAgentRuntimeLockForOwnedMutation(ctx, dbsqlc.LockAgentRuntimeLockForOwnedMutationParams{ + ProjectID: h.route.ProjectID, AgentID: owned.AgentID, ID: owned.RuntimeLockID, + }); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return false, storeerr.ErrRuntimeLockInactive + } + return false, err + } + u.runtimeLocked = true + } + return true, nil +} + +func (u *Unit) LockAgentRefs(ctx context.Context, refs []lifecyclelock.AgentRef, authority MutationAuthority) error { + plan, err := u.PlanAgents(ctx, refs, authority) + if err != nil { + return err + } + return u.LockAgents(ctx, plan) +} + +func (u *Unit) LockProjectDeletion(ctx context.Context, orgID, projectID uuid.UUID) error { + if len(u.handles) != 0 { + return ErrLockPlan + } + if err := lifecyclelock.OrganizationShared(ctx, u.DB(), orgID); err != nil { + return err + } + if err := lifecyclelock.ProjectExclusive(ctx, u.DB(), projectID); err != nil { + return err + } + u.exclusiveProjects[projectID] = struct{}{} + return nil +} + +func (u *Unit) LockOrganizationDeletion(ctx context.Context, orgID uuid.UUID) error { + if len(u.handles) != 0 { + return ErrLockPlan + } + if err := lifecyclelock.OrganizationExclusive(ctx, u.DB(), orgID); err != nil { + return err + } + ids, err := dbsqlc.New(u.DB()).ListActiveProjectIDsForOrganization(ctx, + dbsqlc.ListActiveProjectIDsForOrganizationParams{OrgID: orgID}) + if err != nil { + return err + } + for _, id := range ids { + u.exclusiveProjects[id] = struct{}{} + } + return nil +} + +func (u *Unit) ClaimAgent(ctx context.Context) (AgentRoute, error) { + if len(u.handles) != 0 { + return AgentRoute{}, ErrLockPlan + } + wakeup, err := dbsqlc.New(u.DB()).ClaimNextAgentWakeup(ctx) + if err != nil { + return AgentRoute{}, err + } + plan, err := u.PlanAgents(ctx, []lifecyclelock.AgentRef{{ProjectID: wakeup.ProjectID, AgentID: wakeup.AgentID}}, + IngressAuthority{}) + if err != nil { + return AgentRoute{}, err + } + route := plan.agents[0] + u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} + return route, nil +} + +func (u *Unit) CreatedAgent(projectID, agentID, parentID uuid.UUID) error { + if err := u.active(); err != nil { + return err + } + if projectID == uuid.Nil || agentID == uuid.Nil || u.handles[agentID] != nil { + return ErrLockPlan + } + rootID := agentID + if parentID != uuid.Nil { + parent := u.handles[parentID] + if parent == nil || parent.route.ProjectID != projectID { + return ErrLockPlan + } + rootID = parent.route.RootAgentID + } + route := AgentRoute{CellID: u.cell.id, ProjectID: projectID, AgentID: agentID, RootAgentID: rootID} + u.handles[agentID] = &Handle{unit: u, route: route, valid: true, created: true} + return nil +} + +func (u *Unit) LockAgent( + ctx context.Context, arg dbsqlc.LockAgentInProjectParams, authority MutationAuthority, +) (dbsqlc.LockAgentInProjectRow, error) { + if err := u.LockAgentRefs( + ctx, + []lifecyclelock.AgentRef{{ProjectID: arg.ProjectID, AgentID: arg.ID}}, + authority, + ); err != nil { + if errors.Is(err, storeerr.ErrNotFound) { + err = pgx.ErrNoRows + } + return dbsqlc.LockAgentInProjectRow{}, err + } + return dbsqlc.New(u.DB()).LockAgentInProject(ctx, arg) +} diff --git a/internal/storage/internal/agentexecution/mutation.go b/internal/storage/internal/agentexecution/mutation.go new file mode 100644 index 000000000..51f40e7a5 --- /dev/null +++ b/internal/storage/internal/agentexecution/mutation.go @@ -0,0 +1,159 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" +) + +type executionMutation struct { + head ExecutionHead + insertHead bool + dirty bool + loaded *ExecutionSnapshot +} + +func (m *executionMutation) clone() *executionMutation { + cloned := *m + cloned.loaded = nil + if m.head.LogicalReadyAt != nil { + ready := *m.head.LogicalReadyAt + cloned.head.LogicalReadyAt = &ready + } + return &cloned +} + +func (h *Handle) executionMutation(ctx context.Context) (*executionMutation, error) { + if _, err := h.Route(); err != nil { + return nil, err + } + if h.mutation != nil { + return h.mutation, nil + } + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return nil, err + } + m := &executionMutation{head: snapshot.Head, loaded: &snapshot} + h.mutation = m + return m, nil +} + +func (m *executionMutation) publicationParams(route AgentRoute) executiondb.PublishExecutionParams { + h := m.head + return executiondb.PublishExecutionParams{ + AgentID: route.AgentID, + InsertHead: m.insertHead, + CurrentTurnID: nullableID(h.CurrentTurnID), + StopSequence: h.StopSequence, + AnsweredThroughSequence: h.AnsweredThroughSequence, + MaxNormalInputSequence: h.MaxNormalInputSequence, + MaxContextInputSequence: h.MaxContextInputSequence, + NormalContextID: nullableID(h.NormalContextID), + CompactionContextID: nullableID(h.CompactionContextID), + PendingToolOutputID: nullableID(h.PendingToolOutputID), + PendingOutputLimitID: nullableID(h.PendingOutputLimitID), + PendingConfigInputID: nullableID(h.PendingConfigInputID), + PendingCheckpointID: nullableID(h.PendingCheckpointID), + TurnContinuable: h.TurnContinuable, + IncompleteTools: h.IncompleteTools, + LogicalReadyAt: h.LogicalReadyAt, + } +} + +func (u *Unit) publishExecution(ctx context.Context) error { + var handles []*Handle + var loads []executiondb.LoadExecutionFactsParams + for _, h := range u.orderedHandles() { + if h.mutation == nil || !h.mutation.dirty { + continue + } + handles = append(handles, h) + args := executionFactsParams(h.route, h.mutation.head) + _, captured := u.webhooks[h.route.AgentID] + args.CaptureWebhook = !captured + loads = append(loads, args) + } + if len(handles) == 0 { + return nil + } + var failure error + q := executiondb.New() + batch := q.LoadExecutionFacts(ctx, u.DB(), loads) + batch.Query(func(i int, rows []executiondb.LoadExecutionFactsRow, err error) { + if failure != nil { + return + } + if err != nil { + failure = err + return + } + h := handles[i] + if loads[i].CaptureWebhook { + for _, row := range rows { + if row.Kind != "scope" { + continue + } + target := webhookTarget{orgID: row.WebhookOrgID} + if err := json.Unmarshal(row.WebhookEvents, &target.events); err != nil { + failure = err + return + } + u.webhooks[h.route.AgentID] = target + } + } + snapshot, err := executionSnapshot(h.route, h.mutation.head, rows) + if err != nil { + failure = err + return + } + h.mutation.head = snapshot.Head + }) + if err := errors.Join(failure, batch.Close()); err != nil { + return err + } + writes := make([]executiondb.PublishExecutionParams, len(handles)) + for i, h := range handles { + writes[i] = h.mutation.publicationParams(h.route) + } + published := q.PublishExecution(ctx, u.DB(), writes) + published.QueryRow(func(i int, id uuid.UUID, err error) { + if failure != nil { + return + } + if err != nil { + failure = err + return + } + if id != handles[i].route.AgentID { + failure = fmt.Errorf("%w: missing execution head", ErrInvalidState) + } + }) + return errors.Join(failure, published.Close()) +} + +func executeCommand[T any]( + ctx context.Context, + h *Handle, + run func(*executionMutation) (T, error), +) (result T, err error) { + m, err := h.executionMutation(ctx) + if err != nil { + return result, err + } + defer func() { + if err != nil { + h.unit.aborted = true + } + }() + return run(m) +} + +func (m *executionMutation) changed() { + m.dirty = true + m.loaded = nil +} diff --git a/internal/storage/internal/agentexecution/payload_plans_integration_test.go b/internal/storage/internal/agentexecution/payload_plans_integration_test.go new file mode 100644 index 000000000..1b808a9c6 --- /dev/null +++ b/internal/storage/internal/agentexecution/payload_plans_integration_test.go @@ -0,0 +1,32 @@ +//go:build integration + +package agentexecution_test + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestExecutionPayloadInsertsUseGenericPlans(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + for range 20 { + receiveCommand(t, h, "steering", "payload") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + } + for _, name := range []string{"InsertExecutionContent", "AppendExecutionEvent"} { + var generic, custom int64 + require.NoError( + t, + u.DB(). + QueryRow(t.Context(), + `SELECT generic_plans,custom_plans FROM pg_prepared_statements WHERE statement LIKE $1`, + "-- name: "+name+" :%"). + Scan(&generic, &custom), + ) + require.GreaterOrEqual(t, generic, int64(15), name) + require.LessOrEqual(t, custom, int64(5), name) + } +} diff --git a/internal/storage/internal/agentexecution/publication_integration_test.go b/internal/storage/internal/agentexecution/publication_integration_test.go new file mode 100644 index 000000000..cb13317d0 --- /dev/null +++ b/internal/storage/internal/agentexecution/publication_integration_test.go @@ -0,0 +1,144 @@ +//go:build integration + +package agentexecution_test + +import ( + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/stretchr/testify/require" +) + +func TestExecutionCommitRoundTrips(t *testing.T) { + for _, step := range []string{"model", "tool"} { + t.Run(step, func(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + if step == "model" { + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + _, err = h.AcceptOutput( + t.Context(), + agentexecution.AcceptOutputInput{ + RuntimeLockID: lease, + ContextID: prepared.Context.ID, + Response: commandResponse(modelenvelope.StopReasonMaxTokens), + }, + ) + require.NoError(t, err) + } else { + ref := toolCommand(t, h, lease, "built_in") + _, err := h.AuthorizeTool(t.Context(), ref) + require.NoError(t, err) + _, + err = h.CompleteTool(t.Context(), + agentexecution.ToolCompletion{ToolRef: ref, + Outcome: "succeeded", + Content: []agentexecution.Content{{Kind: "text", + Text: "done"}}}) + require.NoError(t, err) + } + statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) + require.NoError(t, err) + require.Equal(t, 2, statements) + require.Equal(t, 2, roundTrips) + }) + } + for _, size := range []int{2, 1000} { + t.Run(map[int]string{2: "parent_child", 1000: "wide"}[size], func(t *testing.T) { + f := migratedExecutionFixture(t) + root := f.create(t, uuid.New(), uuid.Nil) + ids := []uuid.UUID{root} + for i := 1; i < size; i++ { + id := uuid.New() + f.create(t, id, root) + ids = append(ids, id) + } + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(t.Context()) + require.NoError(t, err) + t.Cleanup(func() { _ = u.Rollback(t.Context()) }) + refs := make([]lifecyclelock.AgentRef, len(ids)) + for i, id := range ids { + refs[i] = lifecyclelock.AgentRef{ProjectID: f.project, AgentID: id} + } + require.NoError(t, u.LockAgentRefs(t.Context(), refs, agentexecution.LifecycleAuthority{})) + for _, id := range ids { + h, err := u.Agent( + agentexecution.AgentRoute{ + CellID: "test", + ProjectID: f.project, + AgentID: id, + RootAgentID: root, + }, + ) + require.NoError(t, err) + receiveCommand(t, h, "queued", "input") + } + statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) + require.NoError(t, err) + require.Equal(t, 2*size, statements) + require.Equal(t, 2, roundTrips) + var heads, wakeups int + require.NoError( + t, + f.pool.QueryRow( + t.Context(), + `SELECT (SELECT count(*) FROM agent_execution_state WHERE logical_ready_at IS NOT NULL),(SELECT count(*) FROM agent_wakeups)`, + ). + Scan(&heads, &wakeups), + ) + require.Equal(t, size, heads) + require.Equal(t, size, wakeups) + + u, err = cell.Begin(t.Context()) + require.NoError(t, err) + t.Cleanup(func() { _ = u.Rollback(t.Context()) }) + require.NoError(t, u.LockAgentRefs(t.Context(), refs, agentexecution.LifecycleAuthority{})) + routes := make([]agentexecution.AgentRoute, len(ids)) + for i, id := range ids { + routes[i] = agentexecution.AgentRoute{ + CellID: "test", + ProjectID: f.project, + AgentID: id, + RootAgentID: root, + } + } + require.NoError(t, u.ArchiveAgents(t.Context(), routes, uuid.Nil)) + statements, roundTrips, err = agentexecution.CountCommit(t.Context(), u) + require.NoError(t, err) + require.Equal(t, 2*size, statements) + require.Equal(t, 2, roundTrips) + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_execution_state +WHERE logical_ready_at IS NOT NULL OR turn_continuable OR incomplete_tools`).Scan(&heads)) + require.Zero(t, heads) + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_wakeups`).Scan(&wakeups), + ) + require.Zero(t, wakeups) + }) + } +} + +func TestExecutionPublicationFailureRollsBack(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + input := receiveCommand(t, h, "queued", "must roll back") + _, err := u.DB().Exec(t.Context(), `DELETE FROM agent_execution_state WHERE agent_id=$1`, a) + require.NoError(t, err) + require.Error(t, u.Commit(t.Context(), "missing head")) + require.NoError(t, u.Rollback(t.Context())) + var exists bool + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT EXISTS(SELECT 1 FROM agent_inputs WHERE id=$1)`, input.ID). + Scan(&exists), + ) + require.False(t, exists) +} diff --git a/internal/storage/internal/agentexecution/queries/commands_agents.sql b/internal/storage/internal/agentexecution/queries/commands_agents.sql new file mode 100644 index 000000000..b14d31569 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_agents.sql @@ -0,0 +1,36 @@ +-- name: CreateExecutionAgent :one +INSERT INTO agents(id,org_id,project_id,root_agent_id,parent_agent_id,subagent_key,state,name,agent_profile_id, + current_config_id,idempotency_key,archive_after_idle_minutes,created_at,updated_at) + SELECT sqlc.arg(id),p.org_id,p.id,sqlc.arg(root_id),sqlc.narg(parent_id),sqlc.arg(subagent_key),'active',sqlc.arg(name), + sqlc.narg(profile_id),sqlc.arg(config_id),sqlc.narg(idempotency_key),sqlc.narg(archive_after_idle_minutes), + statement_timestamp(),statement_timestamp() + FROM projects p JOIN orgs o ON o.id=p.org_id WHERE p.id=sqlc.arg(project_id) AND p.deleted_at IS NULL AND o.deleted_at IS NULL + RETURNING id; + +-- name: ReadExecutionAgentIdentity :one +SELECT org_id,parent_agent_id,root_agent_id,state,name,subagent_key FROM agents +WHERE project_id=sqlc.arg(project_id) AND id=sqlc.arg(id); + +-- name: ArchiveExecutionAgent :execrows +UPDATE agents SET state='archived',archived_at=coalesce(archived_at,statement_timestamp()),updated_at=statement_timestamp() +WHERE id=sqlc.arg(id) AND state='active'; + +-- name: WriteExecutionInteractionTarget :execrows +UPDATE agents a SET interaction_target_id=sqlc.narg(target_id),interaction_handler_key=sqlc.narg(handler_key), + interaction_auto_select=sqlc.arg(auto_select),updated_at=statement_timestamp() +WHERE a.id=sqlc.arg(id) AND ((sqlc.narg(target_id)::uuid IS NULL AND sqlc.narg(handler_key)::text IS NULL) + OR (sqlc.narg(handler_key)::text<>'' AND EXISTS(SELECT 1 FROM integration_targets t JOIN integrations i + ON i.project_id=t.project_id AND i.id=t.integration_id WHERE t.project_id=a.project_id AND t.agent_id=a.id + AND t.id=sqlc.narg(target_id) AND t.deleted_at IS NULL AND i.deleted_at IS NULL AND i.state='active'))); + +-- name: ListExecutionIntegrationTargetAgents :many +-- @sqlc-vet-disable integration-targets-deleted-at +SELECT a.id FROM agents a JOIN integration_targets t ON t.project_id=a.project_id AND t.agent_id=a.id AND t.id=a.interaction_target_id +WHERE a.project_id=sqlc.arg(project_id) AND t.integration_id=sqlc.arg(integration_id) ORDER BY a.id; + +-- name: ClearExecutionIntegrationTargets :exec +-- @sqlc-vet-disable integration-targets-deleted-at +UPDATE agents a SET interaction_target_id=NULL,interaction_handler_key=NULL,updated_at=statement_timestamp() +FROM integration_targets t WHERE a.project_id=sqlc.arg(project_id) AND a.id=ANY(sqlc.arg(agent_ids)::uuid[]) + AND t.project_id=a.project_id AND t.agent_id=a.id AND t.id=a.interaction_target_id AND t.integration_id=sqlc.arg(integration_id); + diff --git a/internal/storage/internal/agentexecution/queries/commands_events.sql b/internal/storage/internal/agentexecution/queries/commands_events.sql new file mode 100644 index 000000000..ca57afc9d --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_events.sql @@ -0,0 +1,45 @@ +-- name: AppendExecutionEvent :one +WITH payload AS MATERIALIZED ( + SELECT agent_payload_created_at(sqlc.arg(project_id),sqlc.arg(agent_id),CASE WHEN sqlc.narg(result_id)::uuid IS NOT NULL THEN 'tool_call_result' ELSE sqlc.arg(kind) END, + coalesce(sqlc.narg(input_id)::uuid,sqlc.narg(output_id)::uuid,sqlc.narg(checkpoint_id)::uuid,sqlc.narg(result_id)::uuid),true) AS created_at +), allocated AS ( + UPDATE agents SET next_event_sequence=next_event_sequence+1 WHERE id=sqlc.arg(agent_id) + RETURNING next_event_sequence-1 AS sequence +) +INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,agent_input_id,model_output_id, + context_checkpoint_id,tool_call_result_id,is_opening_event,created_at,idempotency_key) +SELECT sqlc.arg(id),sqlc.arg(agent_id),sqlc.narg(turn_id),sequence,sqlc.arg(kind),sqlc.narg(input_id), + sqlc.narg(output_id),sqlc.narg(checkpoint_id),sqlc.narg(result_id),sqlc.arg(opening),payload.created_at, + sqlc.narg(idempotency_key) FROM allocated CROSS JOIN payload +RETURNING id,sequence,created_at; + +-- name: OpenExecutionTurn :exec +INSERT INTO agent_turns(id,agent_id,turn_sequence,latest_event_id,latest_semantic_event_id) +SELECT sqlc.arg(id),sqlc.arg(agent_id),coalesce((SELECT max(turn_sequence)+1 FROM agent_turns + WHERE agent_id=sqlc.arg(agent_id)),1),sqlc.arg(event_id),sqlc.arg(event_id); + +-- name: AdvanceExecutionTurn :execrows +UPDATE agent_turns SET latest_event_id=sqlc.arg(event_id), + latest_semantic_event_id=CASE WHEN sqlc.arg(semantic)::boolean THEN sqlc.arg(event_id) ELSE latest_semantic_event_id END +WHERE agent_id=sqlc.arg(agent_id) AND id=sqlc.arg(turn_id); + +-- name: InsertExecutionContent :exec +WITH payload AS MATERIALIZED ( + SELECT agent_payload_created_at(sqlc.arg(project_id),sqlc.arg(agent_id),sqlc.arg(owner_kind), + coalesce(sqlc.narg(input_id)::uuid,sqlc.narg(output_id)::uuid,sqlc.narg(result_id)::uuid),false) AS created_at +) +INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,owner_model_output_id,owner_tool_call_result_id, + ordinal,block_kind,text_content,structured_data,artifact_id,tool_call_id,exclude_from_model_context,metadata,created_at) +SELECT sqlc.arg(agent_id),sqlc.arg(owner_kind),sqlc.narg(input_id),sqlc.narg(output_id),sqlc.narg(result_id), + sqlc.arg(ordinal),sqlc.arg(kind),sqlc.narg(text),sqlc.narg(data)::jsonb,sqlc.narg(artifact_id),sqlc.narg(tool_id), + sqlc.arg(exclude),sqlc.arg(metadata)::jsonb,payload.created_at FROM payload; + +-- name: ReadExecutionContent :many +SELECT ordinal,block_kind,text_content,structured_data,artifact_id,tool_call_id,exclude_from_model_context,metadata +FROM content_blocks WHERE agent_id=sqlc.arg(agent_id) AND + (owner_agent_input_id=sqlc.narg(input_id) OR owner_model_output_id=sqlc.narg(output_id) + OR owner_tool_call_result_id=sqlc.narg(result_id)) ORDER BY ordinal; + +-- name: FindExecutionInputEvent :one +SELECT e.id,e.turn_id,e.sequence,e.created_at FROM agent_events e WHERE e.agent_id=sqlc.arg(agent_id) + AND e.agent_input_id=sqlc.arg(input_id); diff --git a/internal/storage/internal/agentexecution/queries/commands_inputs.sql b/internal/storage/internal/agentexecution/queries/commands_inputs.sql new file mode 100644 index 000000000..f9b883f19 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_inputs.sql @@ -0,0 +1,40 @@ +-- name: FindExecutionInput :one +SELECT id,state,input_kind,delivery_mode,actor_id,integration_target_id,agent_config_id, + metadata,queued_at,admitted_event_id,input_rank +FROM agent_inputs WHERE agent_id=sqlc.arg(agent_id) AND + (id=sqlc.arg(id) OR (idempotency_scope=sqlc.narg(scope) AND input_idempotency_key=sqlc.narg(key))); + +-- name: ReceiveExecutionInput :one +INSERT INTO agent_inputs(id,project_id,agent_id,state,input_rank,actor_id,input_kind,delivery_mode, + integration_target_id,agent_config_id,idempotency_scope,input_idempotency_key,queued_at,metadata) +SELECT sqlc.arg(id),a.project_id,a.id,'received',coalesce((SELECT max(i.input_rank)+1024 + FROM agent_inputs i WHERE i.agent_id=a.id AND i.state='received' AND i.delivery_mode=sqlc.arg(mode)),1024), + sqlc.narg(actor_id),sqlc.arg(kind),sqlc.arg(mode),sqlc.narg(target_id),sqlc.narg(config_id), + sqlc.narg(scope),sqlc.narg(key),statement_timestamp(),sqlc.arg(metadata)::jsonb +FROM agents a WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) AND a.state='active' +RETURNING id,queued_at; + +-- name: ListExecutionBacklog :many +SELECT id,delivery_mode,input_rank,queued_at,integration_target_id,actor_id FROM agent_inputs +WHERE agent_id=sqlc.arg(agent_id) AND input_kind='content' AND state='received' +ORDER BY input_rank,queued_at,id; + +-- name: ChangeExecutionInput :execrows +UPDATE agent_inputs SET delivery_mode=sqlc.arg(mode),input_rank=sqlc.arg(rank), + state=CASE WHEN sqlc.arg(cancel)::boolean THEN 'canceled' ELSE state END, + canceled_at=CASE WHEN sqlc.arg(cancel)::boolean THEN statement_timestamp() ELSE NULL END +WHERE agent_id=sqlc.arg(agent_id) AND id=sqlc.arg(id) AND state='received' AND input_kind='content' + AND (delivery_mode,input_rank,state) IS DISTINCT FROM + (sqlc.arg(mode),sqlc.arg(rank),CASE WHEN sqlc.arg(cancel)::boolean THEN 'canceled' ELSE state END); + +-- name: ResolveExecutionInput :execrows +UPDATE agent_inputs i SET state='resolved',admitted_event_id=sqlc.arg(event_id), + admitted_at=e.created_at,resolved_at=e.created_at, + opening_input_ids=sqlc.narg(opening_ids)::uuid[],opening_event_sequence=sqlc.narg(opening_sequence)::bigint +FROM agent_events e WHERE i.agent_id=sqlc.arg(agent_id) AND i.id=sqlc.arg(id) AND i.state='received' + AND e.agent_id=i.agent_id AND e.id=sqlc.arg(event_id) AND e.agent_input_id=i.id; + +-- name: ActivateExecutionConfig :execrows +UPDATE agents a SET current_config_id=c.id,updated_at=statement_timestamp() +FROM agent_configs c WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) + AND c.project_id=a.project_id AND c.id=sqlc.arg(config_id); diff --git a/internal/storage/internal/agentexecution/queries/commands_interactions.sql b/internal/storage/internal/agentexecution/queries/commands_interactions.sql new file mode 100644 index 000000000..dd04ee1d3 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_interactions.sql @@ -0,0 +1,34 @@ +-- name: ReadExecutionInteraction :one +SELECT i.id,i.tool_call_id,i.interaction_kind,i.state,i.request,i.resolution,i.resolved_by_input_id, + i.destination,i.presentation_receipt +FROM agent_interactions i WHERE i.agent_id=sqlc.arg(agent_id) AND + (i.id=sqlc.arg(id) OR (i.tool_call_id=sqlc.narg(tool_id) AND i.interaction_kind=sqlc.narg(kind))); + +-- name: InsertExecutionInteraction :one +INSERT INTO agent_interactions(agent_id,tool_call_id,interaction_kind,state,request,destination,created_at) +VALUES(sqlc.arg(agent_id),sqlc.arg(tool_id),sqlc.arg(kind),'open',sqlc.arg(request)::jsonb, + sqlc.narg(destination)::jsonb,statement_timestamp()) RETURNING id; + +-- name: ResolveExecutionInteraction :execrows +UPDATE agent_interactions SET state=sqlc.arg(state),resolution=sqlc.arg(resolution)::jsonb, + resolved_by_input_id=sqlc.narg(input_id),resolved_at=statement_timestamp() +WHERE agent_id=sqlc.arg(agent_id) AND id=sqlc.arg(id) AND state='open'; + +-- name: ListExecutionInteractions :many +SELECT i.id FROM agent_interactions i JOIN tool_calls c ON c.agent_id=i.agent_id AND c.id=i.tool_call_id +JOIN model_outputs o ON o.agent_id=c.agent_id AND o.id=c.model_output_id +JOIN model_call_contexts m ON m.agent_id=o.agent_id AND m.id=o.model_call_context_id +WHERE i.agent_id=sqlc.arg(agent_id) AND i.state='open' AND m.turn_id=sqlc.arg(turn_id) +ORDER BY i.id; + +-- name: InsertExecutionResponse :one +INSERT INTO agent_inputs(project_id,agent_id,state,actor_id,input_kind,delivery_mode,target_interaction_id, + idempotency_scope,input_idempotency_key,queued_at,metadata) +VALUES(sqlc.arg(project_id),sqlc.arg(agent_id),'received',sqlc.narg(actor_id),'interaction_response','immediate', + sqlc.arg(interaction_id),'agent_interaction_response',sqlc.arg(interaction_id)::uuid::text,statement_timestamp(),'{}'::jsonb) +RETURNING id; + +-- name: WriteExecutionPresentation :execrows +UPDATE agent_interactions SET presentation_receipt=sqlc.arg(receipt)::jsonb +WHERE agent_id=sqlc.arg(agent_id) AND id=sqlc.arg(id) AND destination=sqlc.arg(destination)::jsonb + AND presentation_receipt IS NULL; diff --git a/internal/storage/internal/agentexecution/queries/commands_models.sql b/internal/storage/internal/agentexecution/queries/commands_models.sql new file mode 100644 index 000000000..7fd16fd91 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_models.sql @@ -0,0 +1,69 @@ +-- name: FenceExecutionRuntime :one +SELECT r.id FROM agent_runtime_locks r JOIN agents a ON a.id=r.agent_id +WHERE a.project_id=sqlc.arg(project_id) AND a.id=sqlc.arg(agent_id) AND a.state='active' + AND r.id=sqlc.arg(runtime_id) AND r.cancel_requested_at IS NULL AND r.lease_expires_at>statement_timestamp(); + +-- name: ExecutionPreparationBoundary :one +SELECT a.current_config_id,(a.next_event_sequence-1)::bigint AS watermark,statement_timestamp()::timestamptz AS database_now, + EXISTS(SELECT 1 FROM model_call_contexts c WHERE c.agent_id=a.id AND c.state='started') AS started +FROM agents a WHERE a.id=sqlc.arg(agent_id) AND a.state='active'; + +-- name: ReadExecutionAttempt :one +SELECT c.id,c.agent_id,c.turn_id,c.operation_kind,c.attempt_number,c.input_event_sequence,c.source_event_sequence_end, + c.state,c.recovery_kind,c.retry_at,c.opening_input_ids,c.opening_event_sequence,c.runtime_lock_id, + c.agent_config_id,c.configured_model_revision_id,c.api_format,c.api_variant,c.provider_request_id,c.provider_response_id, + c.error_kind,c.error_code,c.error_message,c.error_details,c.input_tokens_total,c.uncached_input_tokens, + c.cache_read_input_tokens,c.cache_write_input_tokens,c.output_tokens_total,c.reasoning_output_tokens, + coalesce(c.provider_reported_cost_usd::text,'')::text AS provider_reported_cost_usd,c.provider_metadata, + r.provider_model_slug +FROM model_call_contexts c JOIN configured_model_revisions r ON r.org_id=c.org_id AND r.id=c.configured_model_revision_id +WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=sqlc.arg(id); + +-- name: ExecutionModelRevision :one +-- @sqlc-vet-disable configured-models-deleted-at +-- @sqlc-vet-disable model-provider-configs-deleted-at +SELECT m.current_revision_id,CASE WHEN p.management_kind='cluster' THEN coalesce(w.new_managed_work_allowed,true) + ELSE true END::boolean AS allowed +FROM agent_configs c JOIN configured_models m ON m.org_id=c.org_id AND m.id=c.configured_model_id +JOIN model_provider_configs p ON p.org_id=m.org_id AND p.id=m.model_provider_config_id +LEFT JOIN org_managed_work_admission w ON w.org_id=c.org_id +WHERE c.project_id=sqlc.arg(project_id) AND c.id=sqlc.arg(config_id); + +-- name: CreateExecutionAttempt :one +INSERT INTO model_call_contexts(org_id,project_id,agent_id,id,turn_id,operation_kind,attempt_number, + agent_config_id,configured_model_revision_id,input_event_sequence,source_event_sequence_end,runtime_lock_id, + opening_input_ids,opening_event_sequence,state,created_at) +SELECT a.org_id,a.project_id,a.id,sqlc.arg(id),sqlc.arg(turn_id),sqlc.arg(operation),sqlc.arg(attempt), + sqlc.arg(config_id),sqlc.arg(revision_id),sqlc.arg(watermark),sqlc.narg(source_end),sqlc.arg(runtime_id), + sqlc.arg(opening_ids)::uuid[],sqlc.arg(opening_sequence),'started',statement_timestamp() +FROM agents a WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) + AND EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id AND r.id=sqlc.arg(runtime_id) + AND r.cancel_requested_at IS NULL AND r.lease_expires_at>statement_timestamp()) +RETURNING id; + +-- name: FinishExecutionAttempt :execrows +UPDATE model_call_contexts c SET state=sqlc.arg(state),recovery_kind=sqlc.narg(recovery), + retry_at=CASE WHEN sqlc.narg(retry_microseconds)::bigint IS NULL THEN NULL + ELSE statement_timestamp()+sqlc.narg(retry_microseconds)::bigint*interval '1 microsecond' END, + api_format=sqlc.arg(api_format),api_variant=sqlc.arg(api_variant),provider_request_id=sqlc.arg(request_id), + provider_response_id=sqlc.arg(response_id),error_kind=sqlc.arg(error_kind),error_code=sqlc.arg(error_code), + error_message=sqlc.arg(error_message),error_details=sqlc.arg(error_details)::jsonb, + input_tokens_total=sqlc.narg(input_tokens)::integer,uncached_input_tokens=sqlc.narg(uncached_tokens)::integer, + cache_read_input_tokens=sqlc.narg(cache_read_tokens)::integer,cache_write_input_tokens=sqlc.narg(cache_write_tokens)::integer, + output_tokens_total=sqlc.narg(output_tokens)::integer,reasoning_output_tokens=sqlc.narg(reasoning_tokens)::integer, + provider_reported_cost_usd=sqlc.narg(cost)::text::numeric,provider_metadata=sqlc.arg(provider_metadata)::jsonb, + completed_at=statement_timestamp() +WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=sqlc.arg(id) AND c.state='started' AND c.runtime_lock_id=sqlc.arg(runtime_id) + AND EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=c.agent_id + AND r.id=sqlc.arg(runtime_id) AND r.cancel_requested_at IS NULL AND r.lease_expires_at>statement_timestamp()); + +-- name: LatestExecutionAttempt :one +SELECT id FROM model_call_contexts WHERE agent_id=sqlc.arg(agent_id) AND operation_kind=sqlc.arg(operation) + AND input_event_sequence=sqlc.arg(watermark) + AND source_event_sequence_end IS NOT DISTINCT FROM sqlc.narg(source_end)::bigint +ORDER BY attempt_number DESC LIMIT 1; + +-- name: ExecutionCompactionBoundary :one +SELECT coalesce((SELECT summarized_through_event_sequence FROM context_checkpoints + WHERE agent_id=sqlc.arg(agent_id) AND summarized_through_event_sequence<=sqlc.arg(watermark) + ORDER BY summarized_through_event_sequence DESC LIMIT 1),0)::bigint AS summarized_through; diff --git a/internal/storage/internal/agentexecution/queries/commands_outputs.sql b/internal/storage/internal/agentexecution/queries/commands_outputs.sql new file mode 100644 index 000000000..a2ac60eb5 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_outputs.sql @@ -0,0 +1,40 @@ +-- name: CreateExecutionOutput :one +INSERT INTO model_outputs(agent_id,model_call_context_id,served_provider_model_slug,stop_reason,provider_replay,created_at) +VALUES(sqlc.arg(agent_id),sqlc.arg(context_id),sqlc.arg(served_model),sqlc.arg(stop_reason),sqlc.narg(replay)::jsonb,statement_timestamp()) +RETURNING id,created_at; + +-- name: FindExecutionOutput :one +SELECT o.id,o.served_provider_model_slug,o.stop_reason,o.provider_replay,e.id AS event_id,e.turn_id,e.sequence,e.created_at +FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +WHERE o.agent_id=sqlc.arg(agent_id) AND o.model_call_context_id=sqlc.arg(context_id); + +-- name: CreateExecutionToolProposal :exec +INSERT INTO tool_calls(id,agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) +SELECT sqlc.arg(id),sqlc.arg(agent_id),o.id,sqlc.arg(provider_call_id),sqlc.arg(name),sqlc.arg(input)::jsonb, + sqlc.arg(type),'awaiting_authorization',o.created_at FROM model_outputs o + WHERE o.agent_id=sqlc.arg(agent_id) AND o.id=sqlc.arg(output_id); + +-- name: ReadExecutionProposals :many +SELECT id,provider_call_id,name,input,type FROM tool_calls +WHERE agent_id=sqlc.arg(agent_id) AND model_output_id=sqlc.arg(output_id) ORDER BY provider_call_id; + +-- name: ValidateExecutionCheckpoint :one +SELECT (SELECT count(*) FROM agent_events e WHERE e.agent_id=sqlc.arg(agent_id) + AND e.sequence BETWEEN sqlc.arg(start_sequence) AND sqlc.arg(end_sequence))::bigint AS events, + EXISTS(SELECT 1 FROM agent_events e JOIN tool_calls c ON c.agent_id=e.agent_id AND c.model_output_id=e.model_output_id + LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id + LEFT JOIN agent_events re ON re.agent_id=r.agent_id AND re.tool_call_result_id=r.id + WHERE e.agent_id=sqlc.arg(agent_id) AND e.sequence BETWEEN sqlc.arg(start_sequence) AND sqlc.arg(end_sequence) + AND (re.id IS NULL OR re.sequence>sqlc.arg(end_sequence))) AS open_authorities; + +-- name: CreateExecutionCheckpoint :one +INSERT INTO context_checkpoints(agent_id,summarized_through_event_sequence,producer_model_call_context_id,summary, + opening_input_ids,opening_event_sequence,created_at) +VALUES(sqlc.arg(agent_id),sqlc.arg(source_end),sqlc.arg(context_id),sqlc.arg(summary), + sqlc.arg(opening_ids)::uuid[],sqlc.narg(opening_sequence),statement_timestamp()) +RETURNING id,created_at; + +-- name: FindExecutionCheckpoint :one +SELECT c.id,c.summary,c.summarized_through_event_sequence,e.id AS event_id,e.turn_id,e.sequence,e.created_at +FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=sqlc.arg(agent_id) AND c.producer_model_call_context_id=sqlc.arg(context_id); diff --git a/internal/storage/internal/agentexecution/queries/commands_processes.sql b/internal/storage/internal/agentexecution/queries/commands_processes.sql new file mode 100644 index 000000000..0438fa6f9 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_processes.sql @@ -0,0 +1,10 @@ +-- name: ReadExecutionProcessOwner :one +SELECT p.id,p.org_id,p.machine_id,p.tool_call_id,p.state,p.state_reason_code,p.state_reason_message,p.exit_code, + p.exit_signal,p.execution_spec,p.execution_granted_at,p.source_started_at,p.source_ended_at,p.default_output_cursor +FROM processes p WHERE p.agent_id=sqlc.arg(agent_id) AND p.id=sqlc.arg(id) FOR UPDATE; + +-- name: ReadExecutionActionOwner :one +SELECT a.id,a.process_id,a.tool_call_id,a.action_kind,a.seq,a.payload,a.state,a.state_reason_code,a.state_reason_message, + EXISTS(SELECT 1 FROM process_actions earlier WHERE earlier.agent_id=a.agent_id AND earlier.process_id=a.process_id + AND earlier.seq'completed' + AND (sqlc.narg(turn_id)::uuid IS NULL OR m.turn_id=sqlc.narg(turn_id)) + AND (sqlc.narg(runtime_id)::uuid IS NULL OR (c.state='running' AND c.runtime_lock_id=sqlc.narg(runtime_id))) ORDER BY c.id; + +-- name: InsertExecutionStop :one +INSERT INTO agent_inputs(project_id,agent_id,state,actor_id,input_kind,delivery_mode,control_type, + idempotency_scope,input_idempotency_key,queued_at,metadata) +SELECT a.project_id,a.id,'received',sqlc.narg(actor_id),'control','immediate','cancel_current','agent_control', + 'agent-cancel:'||a.id::text||':after:'||(a.next_event_sequence-1)::text,statement_timestamp(),'{}'::jsonb +FROM agents a WHERE a.id=sqlc.arg(agent_id) RETURNING id; + +-- name: CancelExecutionInputs :execrows +UPDATE agent_inputs SET state='canceled',canceled_at=statement_timestamp() +WHERE agent_id=sqlc.arg(agent_id) AND state='received' AND input_kind='content' + AND (delivery_mode='steering' OR sqlc.arg(include_queued)::boolean); + +-- name: ConsumeExecutionWakeup :exec +DELETE FROM agent_wakeups WHERE agent_id=sqlc.arg(agent_id); + +-- name: ExecutionRetainedRuntimeWork :one +SELECT EXISTS(SELECT 1 FROM model_call_contexts c WHERE c.agent_id=sqlc.arg(agent_id) + AND c.runtime_lock_id=sqlc.arg(runtime_id) AND c.state='started') OR + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) AND c.runtime_lock_id=sqlc.arg(runtime_id) + AND c.state='running') AS retained; diff --git a/internal/storage/internal/agentexecution/queries/commands_tools.sql b/internal/storage/internal/agentexecution/queries/commands_tools.sql new file mode 100644 index 000000000..582edf6e7 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/commands_tools.sql @@ -0,0 +1,40 @@ +-- name: ReadExecutionTool :one +SELECT c.id,c.agent_id,c.model_output_id,c.provider_call_id,c.name,c.input,c.type,c.state,c.runtime_lock_id, + m.turn_id,m.id AS context_id,e.id AS source_event_id,e.sequence AS source_sequence, + r.id AS result_id,r.outcome,r.completed_at,re.id AS result_event_id,re.sequence AS result_sequence +FROM tool_calls c JOIN model_outputs o ON o.agent_id=c.agent_id AND o.id=c.model_output_id +JOIN model_call_contexts m ON m.agent_id=o.agent_id AND m.id=o.model_call_context_id +JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id +LEFT JOIN agent_events re ON re.agent_id=r.agent_id AND re.tool_call_result_id=r.id +WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=sqlc.arg(id); + +-- name: TransitionExecutionTool :execrows +UPDATE tool_calls c SET state=sqlc.arg(next_state),runtime_lock_id=sqlc.narg(next_runtime_id) +WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=sqlc.arg(id) AND c.state=sqlc.arg(previous_state) + AND c.runtime_lock_id IS NOT DISTINCT FROM sqlc.narg(previous_runtime_id)::uuid + AND (sqlc.narg(fence_id)::uuid IS NULL OR EXISTS(SELECT 1 FROM agent_runtime_locks r + WHERE r.agent_id=c.agent_id AND r.id=sqlc.narg(fence_id) AND r.cancel_requested_at IS NULL + AND r.lease_expires_at>statement_timestamp())); + +-- name: InsertExecutionToolResult :one +INSERT INTO tool_call_results(agent_id,tool_call_id,outcome,completed_at) +SELECT c.agent_id,c.id,sqlc.arg(outcome),statement_timestamp() FROM tool_calls c +WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=sqlc.arg(id) AND c.state='completed' +RETURNING id; + +-- name: CloseExecutionToolInteractions :many +UPDATE agent_interactions SET state='canceled',resolution=jsonb_build_object('reason','tool_call_completed'),resolved_at=statement_timestamp() +WHERE agent_id=sqlc.arg(agent_id) AND tool_call_id=sqlc.arg(tool_id) AND state='open' +RETURNING id,interaction_kind,state; + +-- name: ExecutionToolOwner :one +SELECT EXISTS(SELECT 1 FROM agent_interactions i WHERE i.agent_id=sqlc.arg(agent_id) AND i.tool_call_id=sqlc.arg(tool_id) + AND i.id=sqlc.narg(interaction_id) AND i.interaction_kind='question' AND i.state='open') AS question, + EXISTS(SELECT 1 FROM processes p WHERE p.agent_id=sqlc.arg(agent_id) AND p.tool_call_id=sqlc.arg(tool_id) + AND p.id=sqlc.narg(process_id) AND p.state IN ('queued','starting','running')) AS process, + EXISTS(SELECT 1 FROM process_actions p WHERE p.agent_id=sqlc.arg(agent_id) AND p.tool_call_id=sqlc.arg(tool_id) + AND p.id=sqlc.narg(action_id) AND p.state IN ('queued','accepted')) AS action; + +-- name: ReadExecutionToolDispatch :one +SELECT state FROM tool_calls WHERE agent_id=sqlc.arg(agent_id) AND id=sqlc.arg(id); diff --git a/internal/storage/internal/agentexecution/queries/contexts.sql b/internal/storage/internal/agentexecution/queries/contexts.sql new file mode 100644 index 000000000..05446aba3 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/contexts.sql @@ -0,0 +1,35 @@ +-- name: ReconstructExecutionContexts :many +SELECT c.id,c.agent_id,c.turn_id,c.operation_kind,c.attempt_number,c.input_event_sequence, + c.source_event_sequence_end,c.state,c.recovery_kind,c.retry_at,c.opening_input_ids,c.opening_event_sequence +FROM model_call_contexts c WHERE c.agent_id=$1 +ORDER BY c.input_event_sequence,c.attempt_number; + +-- name: ValidateExecutionLineage :one +WITH snapshots AS ( + SELECT c.opening_input_ids,c.opening_event_sequence,c.input_event_sequence AS watermark + FROM model_call_contexts c WHERE c.agent_id=$1 + UNION ALL + SELECT i.opening_input_ids,i.opening_event_sequence,e.sequence + FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id + WHERE i.agent_id=$1 AND i.input_kind='config_change' AND i.state='resolved' + UNION ALL + SELECT c.opening_input_ids,c.opening_event_sequence,e.sequence + FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id + WHERE c.agent_id=$1 +) +SELECT coalesce(NOT EXISTS ( + SELECT 1 FROM model_call_contexts c WHERE c.agent_id=$1 + AND c.turn_id IS DISTINCT FROM (SELECT e.turn_id FROM agent_events e + WHERE e.agent_id=c.agent_id AND e.is_opening_event AND e.sequence<=c.input_event_sequence + ORDER BY e.sequence DESC LIMIT 1) +) AND NOT EXISTS ( + SELECT 1 FROM snapshots s WHERE EXISTS ( + SELECT 1 FROM ( + SELECT o.ordinality,e.sequence,lag(e.sequence) OVER (ORDER BY o.ordinality) AS previous + FROM unnest(s.opening_input_ids) WITH ORDINALITY o(id,ordinality) + LEFT JOIN agent_inputs i ON i.agent_id=$1 AND i.id=o.id AND i.state='resolved' AND i.input_kind='content' + LEFT JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id AND e.is_opening_event + ) opening WHERE sequence IS NULL OR sequence>s.watermark + OR previous>=sequence OR (ordinality=1 AND sequence<>s.opening_event_sequence) + ) +),false)::boolean AS valid; diff --git a/internal/storage/internal/agentexecution/queries/inputs.sql b/internal/storage/internal/agentexecution/queries/inputs.sql new file mode 100644 index 000000000..d207f3f89 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/inputs.sql @@ -0,0 +1,30 @@ +-- name: CaptureExecutionOpening :many +WITH boundary AS MATERIALIZED ( + SELECT greatest( + coalesce((SELECT output.sequence FROM agent_events output WHERE output.agent_id=sqlc.arg(agent_id) + AND output.event_kind='model_output' AND output.sequence<=sqlc.arg(watermark) + ORDER BY output.sequence DESC LIMIT 1),0), + sqlc.arg(stop_sequence)::bigint + )::bigint AS sequence +), unanswered AS MATERIALIZED ( + SELECT i.id,e.sequence,e.created_at FROM boundary b + JOIN agent_events e ON e.agent_id=sqlc.arg(agent_id) AND e.is_opening_event + AND e.sequence>b.sequence AND e.sequence<=sqlc.arg(watermark) + JOIN agent_inputs i ON i.project_id=sqlc.arg(project_id) AND i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id +) +SELECT id,sequence,created_at FROM unanswered +UNION ALL +SELECT i.id,e.sequence,e.created_at FROM agent_events e +JOIN agent_inputs i ON i.project_id=sqlc.arg(project_id) AND i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id +WHERE e.agent_id=sqlc.arg(agent_id) AND e.turn_id=sqlc.arg(turn_id) AND e.is_opening_event + AND e.sequence<=sqlc.arg(watermark) AND NOT EXISTS(SELECT 1 FROM unanswered) +ORDER BY sequence; + +-- name: ReconstructExecutionConfig :one +SELECT i.id FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id +WHERE i.agent_id=$1 AND e.turn_id=$2 AND i.input_kind='config_change' AND i.state='resolved' + AND e.sequence>sqlc.arg(normal_watermark) AND e.sequence>=sqlc.arg(stop_sequence) + AND e.sequence>=sqlc.arg(first_content_sequence) AND sqlc.arg(first_content_sequence)::bigint>0 +ORDER BY e.sequence DESC LIMIT 1; diff --git a/internal/storage/internal/agentexecution/queries/load.sql b/internal/storage/internal/agentexecution/queries/load.sql new file mode 100644 index 000000000..b7179322b --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/load.sql @@ -0,0 +1,271 @@ +-- name: LoadExecutionBase :one +SELECT coalesce(h.agent_id,a.id)::uuid AS agent_id,h.current_turn_id, + coalesce(h.stop_sequence,0)::bigint AS stop_sequence, + coalesce(h.answered_through_sequence,0)::bigint AS answered_through_sequence, + coalesce(h.max_normal_input_sequence,0)::bigint AS max_normal_input_sequence, + coalesce(h.max_context_input_sequence,0)::bigint AS max_context_input_sequence, + h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id, + h.pending_config_input_id,h.pending_checkpoint_id,coalesce(h.turn_continuable,false)::boolean AS turn_continuable, + coalesce(h.incomplete_tools,false)::boolean AS incomplete_tools,h.logical_ready_at, + (h.agent_id IS NOT NULL)::boolean AS head_exists,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now +FROM agents a LEFT JOIN agent_execution_state h ON h.agent_id=a.id +WHERE a.project_id=sqlc.arg(project_id) AND a.id=sqlc.arg(agent_id); + +-- name: LoadExecutionFacts :batchmany +WITH scope AS MATERIALIZED ( + SELECT a.id,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now,t.id AS turn_id, + coalesce(e.sequence,0)::bigint AS semantic_sequence,coalesce(e.created_at,'epoch'::timestamptz)::timestamptz AS semantic_time, + EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='steering' AND i.state='received') AS steering, + EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='queued' AND i.state='received') AS queued, + coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event + ORDER BY x.sequence LIMIT 1),0)::bigint AS first_opening_sequence, + coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event + ORDER BY x.sequence DESC LIMIT 1),0)::bigint AS last_opening_sequence, + coalesce((SELECT x.sequence FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id + WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event AND i.input_kind='content' + ORDER BY x.sequence LIMIT 1),0)::bigint AS first_content_sequence + FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=sqlc.narg(turn_id) + LEFT JOIN agent_events e ON e.agent_id=a.id AND e.id=t.latest_semantic_event_id + WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) +), outputs AS MATERIALIZED ( + SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at + FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id + WHERE o.agent_id=sqlc.arg(agent_id) AND o.id IN (sqlc.narg(tool_output_id)::uuid,sqlc.narg(limit_output_id)::uuid) +), batch AS MATERIALIZED ( + SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) + AND c.model_output_id=sqlc.narg(tool_output_id)) AS has_tools, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) + AND c.model_output_id=sqlc.narg(tool_output_id) AND c.state<>'completed') AS incomplete, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) + AND c.model_output_id=sqlc.narg(tool_output_id) AND + (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable + WHERE sqlc.narg(tool_output_id)::uuid IS NOT NULL +), boundary AS MATERIALIZED ( + SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=sqlc.arg(agent_id) + AND e.event_kind='model_output' AND e.sequence<=(SELECT last_opening_sequence FROM scope) + ORDER BY e.sequence DESC LIMIT 1),0),sqlc.arg(stop_sequence)::bigint)::bigint AS sequence + FROM scope WHERE scope.state='active' AND sqlc.arg(max_context_input_sequence)::bigintb.sequence AND e.sequence<=(SELECT last_opening_sequence FROM scope) + JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id +), opening AS ( + SELECT id,sequence,created_at FROM unanswered + UNION ALL + SELECT i.id,e.sequence,e.created_at FROM boundary b + JOIN agent_events e ON e.agent_id=sqlc.arg(agent_id) AND e.turn_id=sqlc.narg(turn_id) AND e.is_opening_event + AND e.sequence<=(SELECT last_opening_sequence FROM scope) + JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id + WHERE NOT EXISTS(SELECT 1 FROM unanswered) +), facts AS ( +SELECT 'context'::text AS kind, + c.id AS id, + c.turn_id AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + 0::bigint AS sequence, + 'epoch'::timestamptz AS event_time, + c.operation_kind AS operation, + c.attempt_number AS attempt, + c.input_event_sequence AS watermark, + coalesce(c.source_event_sequence_end,0)::bigint AS source_end, + c.state AS state, + coalesce(c.recovery_kind,'')::text AS recovery, + c.retry_at AS retry_at, + c.opening_input_ids AS opening_ids, + c.opening_event_sequence AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM model_call_contexts c WHERE c.agent_id=sqlc.arg(agent_id) AND c.id IN ( + SELECT sqlc.narg(normal_context_id)::uuid UNION SELECT sqlc.narg(compaction_context_id)::uuid + UNION SELECT model_call_context_id FROM outputs) +UNION ALL +SELECT 'output'::text AS kind, + o.id AS id, + '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, + o.model_call_context_id AS context_id, + o.event_id AS event_id, + o.sequence AS sequence, + o.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + '{}'::uuid[] AS opening_ids, + 0::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + o.stop_reason AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM outputs o +UNION ALL +SELECT 'batch'::text AS kind, + '00000000-0000-0000-0000-000000000000'::uuid AS id, + '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + coalesce(r.last_sequence,0)::bigint AS sequence, + coalesce(r.ready_at,'epoch'::timestamptz)::timestamptz AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + '{}'::uuid[] AS opening_ids, + 0::bigint AS opening_sequence, + b.has_tools AS has_tools, + b.incomplete AS incomplete, + b.runnable AS runnable, + coalesce(r.calls,0)::bigint AS calls, + coalesce(r.results,0)::bigint AS results, + coalesce(r.events,0)::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM batch b LEFT JOIN LATERAL ( + SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, + max(e.sequence)::bigint AS last_sequence,max(r.completed_at)::timestamptz AS ready_at + FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id + LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id + WHERE c.agent_id=sqlc.arg(agent_id) AND c.model_output_id=sqlc.narg(tool_output_id) AND NOT b.incomplete +) r ON NOT b.incomplete +UNION ALL +SELECT 'config'::text AS kind, + i.id AS id, + e.turn_id AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + e.sequence AS sequence, + e.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + i.opening_input_ids AS opening_ids, + coalesce(i.opening_event_sequence,0)::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id +WHERE i.agent_id=sqlc.arg(agent_id) AND i.id=sqlc.narg(config_id) AND i.input_kind='config_change' AND i.state='resolved' +UNION ALL +SELECT 'checkpoint'::text AS kind, + c.id AS id, + e.turn_id AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + e.sequence AS sequence, + c.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + c.opening_input_ids AS opening_ids, + coalesce(c.opening_event_sequence,0)::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=sqlc.narg(checkpoint_id) +UNION ALL +SELECT 'opening'::text AS kind, + o.id AS id, + '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, + '00000000-0000-0000-0000-000000000000'::uuid AS context_id, + '00000000-0000-0000-0000-000000000000'::uuid AS event_id, + o.sequence AS sequence, + o.created_at AS event_time, + ''::text AS operation, + 0::integer AS attempt, + 0::bigint AS watermark, + 0::bigint AS source_end, + ''::text AS state, + ''::text AS recovery, + NULL::timestamptz AS retry_at, + '{}'::uuid[] AS opening_ids, + 0::bigint AS opening_sequence, + false AS has_tools, + false AS incomplete, + false AS runnable, + 0::bigint AS calls, + 0::bigint AS results, + 0::bigint AS events, + ''::text AS stop_reason, + '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, + 'epoch'::timestamptz AS database_now, + false AS steering,false AS queued, + 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence +FROM opening o +UNION ALL +SELECT 'scope'::text AS kind,s.id,coalesce(s.turn_id,'00000000-0000-0000-0000-000000000000'::uuid), + '00000000-0000-0000-0000-000000000000'::uuid,'00000000-0000-0000-0000-000000000000'::uuid, + s.semantic_sequence,s.semantic_time,''::text,0::integer,0::bigint,0::bigint,s.state,''::text,NULL::timestamptz, + '{}'::uuid[],0::bigint,false,false,false,0::bigint,0::bigint,0::bigint,''::text, + s.root_agent_id,s.database_now,s.steering,s.queued,s.first_opening_sequence,s.last_opening_sequence,s.first_content_sequence +FROM scope s + ) +SELECT f.kind,f.id,f.turn_id,f.context_id,f.event_id,f.sequence,f.event_time,f.operation,f.attempt, + f.watermark,f.source_end,f.state,f.recovery,f.retry_at,f.opening_ids,f.opening_sequence, + f.has_tools,f.incomplete,f.runnable,f.calls,f.results,f.events,f.stop_reason,f.root_agent_id, + f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence, + coalesce(w.org_id,'00000000-0000-0000-0000-000000000000'::uuid)::uuid AS webhook_org_id, + coalesce(w.events,'[]'::jsonb)::jsonb AS webhook_events +FROM facts f LEFT JOIN LATERAL ( + SELECT p.org_id,c.compiled_definition->'event_webhook'->'events' AS events + FROM agents a JOIN agent_configs c ON c.project_id=a.project_id AND c.id=a.current_config_id + JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id + WHERE a.id=sqlc.arg(agent_id) AND p.deleted_at IS NULL AND o.deleted_at IS NULL + AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' + AND sqlc.arg(capture_webhook)::boolean AND f.kind='scope' +) w ON true ORDER BY f.kind,f.sequence; diff --git a/internal/storage/internal/agentexecution/queries/outputs.sql b/internal/storage/internal/agentexecution/queries/outputs.sql new file mode 100644 index 000000000..d1066cc51 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/outputs.sql @@ -0,0 +1,5 @@ +-- name: ReconstructExecutionOutputs :many +SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at +FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +WHERE o.agent_id=$1 AND e.turn_id=$2 AND e.sequence>=sqlc.arg(stop_sequence) +ORDER BY e.sequence; diff --git a/internal/storage/internal/agentexecution/queries/publication.sql b/internal/storage/internal/agentexecution/queries/publication.sql new file mode 100644 index 000000000..f321ba741 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/publication.sql @@ -0,0 +1,43 @@ +-- name: PublishExecution :batchone +WITH scope AS MATERIALIZED ( + SELECT a.id AS agent_id,sqlc.narg(logical_ready_at)::timestamptz AS logical_ready_at + FROM agents a WHERE a.id=sqlc.arg(agent_id) AND (sqlc.arg(insert_head)::boolean OR + EXISTS(SELECT 1 FROM agent_execution_state h WHERE h.agent_id=a.id)) +), updated AS ( +INSERT INTO agent_execution_state AS h(agent_id,current_turn_id,stop_sequence,answered_through_sequence,max_normal_input_sequence,max_context_input_sequence,normal_context_id,compaction_context_id,pending_tool_output_id,pending_output_limit_id,pending_config_input_id,pending_checkpoint_id,turn_continuable,incomplete_tools,logical_ready_at) +SELECT scope.agent_id,sqlc.narg(current_turn_id),sqlc.arg(stop_sequence),sqlc.arg(answered_through_sequence),sqlc.arg(max_normal_input_sequence),sqlc.arg(max_context_input_sequence),sqlc.narg(normal_context_id),sqlc.narg(compaction_context_id),sqlc.narg(pending_tool_output_id),sqlc.narg(pending_output_limit_id),sqlc.narg(pending_config_input_id),sqlc.narg(pending_checkpoint_id),sqlc.arg(turn_continuable),sqlc.arg(incomplete_tools),scope.logical_ready_at FROM scope +ON CONFLICT(agent_id) DO UPDATE SET +current_turn_id=excluded.current_turn_id, +stop_sequence=excluded.stop_sequence, +answered_through_sequence=excluded.answered_through_sequence, +max_normal_input_sequence=excluded.max_normal_input_sequence, +max_context_input_sequence=excluded.max_context_input_sequence, +normal_context_id=excluded.normal_context_id, +compaction_context_id=excluded.compaction_context_id, +pending_tool_output_id=excluded.pending_tool_output_id, +pending_output_limit_id=excluded.pending_output_limit_id, +pending_config_input_id=excluded.pending_config_input_id, +pending_checkpoint_id=excluded.pending_checkpoint_id, +turn_continuable=excluded.turn_continuable, +incomplete_tools=excluded.incomplete_tools, +logical_ready_at=excluded.logical_ready_at +WHERE (h.current_turn_id,h.stop_sequence,h.answered_through_sequence,h.max_normal_input_sequence,h.max_context_input_sequence,h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id,h.pending_config_input_id,h.pending_checkpoint_id,h.turn_continuable,h.incomplete_tools,h.logical_ready_at) IS DISTINCT FROM +(excluded.current_turn_id,excluded.stop_sequence,excluded.answered_through_sequence,excluded.max_normal_input_sequence,excluded.max_context_input_sequence,excluded.normal_context_id,excluded.compaction_context_id,excluded.pending_tool_output_id,excluded.pending_output_limit_id,excluded.pending_config_input_id,excluded.pending_checkpoint_id,excluded.turn_continuable,excluded.incomplete_tools,excluded.logical_ready_at) +RETURNING h.agent_id), desired AS MATERIALIZED ( + SELECT a.id,h.logical_ready_at AS ready_at FROM scope h JOIN agents a ON a.id=h.agent_id + WHERE a.id=sqlc.arg(agent_id) AND a.state='active' AND h.logical_ready_at IS NOT NULL + AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id) +), removed AS ( + DELETE FROM agent_wakeups WHERE agent_id=sqlc.arg(agent_id) AND NOT EXISTS(SELECT 1 FROM desired) +), published AS ( +INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) +SELECT id,ready_at,statement_timestamp() FROM desired +ON CONFLICT(agent_id) DO UPDATE SET ready_at=CASE + WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() + THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END, + updated_at=statement_timestamp() +WHERE agent_wakeups.ready_at IS DISTINCT FROM CASE + WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() + THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END +RETURNING agent_id +) SELECT agent_id FROM scope; diff --git a/internal/storage/internal/agentexecution/queries/state.sql b/internal/storage/internal/agentexecution/queries/state.sql new file mode 100644 index 000000000..0688f8aa4 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/state.sql @@ -0,0 +1,13 @@ +-- name: LoadExecutionScope :one +SELECT a.root_agent_id +FROM agents a WHERE a.project_id=$1 AND a.id=$2; + +-- name: LoadExecutionBoundaries :one +SELECT coalesce((SELECT max(e.sequence) FROM agent_events e + JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + WHERE i.project_id=$1 AND i.agent_id=$2 AND i.input_kind='control' + AND i.state='resolved' AND i.control_type='cancel_current'),0)::bigint AS stop_sequence, + coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$2 AND e.event_kind='model_output' + ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS output_sequence, + coalesce((SELECT t.id FROM agent_turns t WHERE t.agent_id=$2 ORDER BY t.turn_sequence DESC LIMIT 1), + '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS turn_id; diff --git a/internal/storage/internal/agentexecution/queries/tools.sql b/internal/storage/internal/agentexecution/queries/tools.sql new file mode 100644 index 000000000..87a9bd04f --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/tools.sql @@ -0,0 +1,13 @@ +-- name: LoadExecutionBatch :one +SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2) AS has_tools, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2 AND c.state<>'completed') AS incomplete, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2 AND + (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable; + +-- name: LoadExecutionBatchCompletion :one +SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, + coalesce(max(e.sequence),0)::bigint AS last_result_sequence, + coalesce(max(r.completed_at),'epoch'::timestamptz)::timestamptz AS ready_at +FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id +LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id +WHERE c.agent_id=$1 AND c.model_output_id=$2; diff --git a/internal/storage/internal/agentexecution/queries/turns_events.sql b/internal/storage/internal/agentexecution/queries/turns_events.sql new file mode 100644 index 000000000..736010150 --- /dev/null +++ b/internal/storage/internal/agentexecution/queries/turns_events.sql @@ -0,0 +1,16 @@ +-- name: LoadExecutionTurn :one +SELECT t.id,semantic.sequence AS semantic_sequence,semantic.created_at AS semantic_time, + coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=t.agent_id AND e.turn_id=t.id AND e.is_opening_event + ORDER BY e.sequence LIMIT 1),0)::bigint AS first_opening_sequence, + coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=t.agent_id AND e.turn_id=t.id AND e.is_opening_event + ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS last_opening_sequence, + coalesce((SELECT e.sequence FROM agent_events e JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id + WHERE e.agent_id=t.agent_id AND e.turn_id=t.id AND e.is_opening_event AND i.input_kind='content' + ORDER BY e.sequence LIMIT 1),0)::bigint AS first_content_sequence +FROM agent_turns t JOIN agent_events semantic ON semantic.agent_id=t.agent_id AND semantic.id=t.latest_semantic_event_id +WHERE t.agent_id=$1 AND t.id=$2; + +-- name: ReconstructExecutionCheckpoint :one +SELECT c.id FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=$1 AND e.turn_id=$2 AND e.sequence>sqlc.arg(normal_watermark) AND e.sequence>=sqlc.arg(stop_sequence) +ORDER BY e.sequence DESC LIMIT 1; diff --git a/internal/storage/internal/agentexecution/reconstruction_integration_test.go b/internal/storage/internal/agentexecution/reconstruction_integration_test.go new file mode 100644 index 000000000..fbdefc980 --- /dev/null +++ b/internal/storage/internal/agentexecution/reconstruction_integration_test.go @@ -0,0 +1,536 @@ +//go:build integration + +package agentexecution_test + +import ( + "context" + "io/fs" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "testing/fstest" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/pgx/v5/stdlib" + "github.com/omnara-ai/omnara/internal/dbmigrate" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + schemamigrations "github.com/omnara-ai/omnara/migrations" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + integrationdb.RunTestMain(m) +} + +type executionFixture struct { + pool *pgxpool.Pool + org uuid.UUID + project uuid.UUID + config uuid.UUID + revision uuid.UUID + now time.Time +} + +func releasedExecutionFixture(t *testing.T) executionFixture { + t.Helper() + ctx := t.Context() + pool := integrationdb.OpenUnmigratedPool(t, ctx) + files := fstest.MapFS{} + entries, err := os.ReadDir("../../../../migrations") + require.NoError(t, err) + for _, entry := range entries { + if filepath.Ext(entry.Name()) != ".sql" { + continue + } + version, parseErr := strconv.Atoi(strings.Split(entry.Name(), "_")[0]) + require.NoError(t, parseErr) + if version > 51 { + continue + } + body, readErr := os.ReadFile(filepath.Join("../../../../migrations", entry.Name())) + require.NoError(t, readErr) + files[entry.Name()] = &fstest.MapFile{Data: body, Mode: fs.FileMode(0600)} + } + db := stdlib.OpenDBFromPool(pool) + require.NoError(t, dbmigrate.ApplyPostgres(ctx, db, files, schemamigrations.GoMigrations()...)) + require.NoError(t, db.Close()) + return seedExecutionFixture(t, pool) +} + +func (f executionFixture) agent(t *testing.T) uuid.UUID { + t.Helper() + id := uuid.New() + _, err := f.pool.Exec( + t.Context(), + `INSERT INTO agents(id,org_id,project_id,state,current_config_id,created_at,updated_at) + VALUES($1,$2,$3,'active',$4,$5,$5)`, + id, + f.org, + f.project, + f.config, + f.now, + ) + require.NoError(t, err) + return id +} + +func (f executionFixture) opening( + t *testing.T, + agent uuid.UUID, + sequence int64, + turn uuid.UUID, + kind string, +) (uuid.UUID, uuid.UUID) { + t.Helper() + ctx := t.Context() + tx, err := f.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + input, event := uuid.New(), uuid.New() + fresh := turn == uuid.Nil + if fresh { + turn = uuid.New() + } + var config *uuid.UUID + mode := "steering" + if kind == "config_change" { + config = &f.config + mode = "immediate" + } + _, err = tx.Exec( + ctx, + `INSERT INTO agent_inputs(id,project_id,agent_id,state,input_kind,delivery_mode,agent_config_id,queued_at) + VALUES($1,$2,$3,'received',$4,$5,$6,$7)`, + input, + f.project, + agent, + kind, + mode, + config, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + ` +INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,agent_input_id,is_opening_event, +idempotency_key,created_at) + VALUES($1,$2,$3,$4,'agent_input',$5,$6,'agent_input:'||($5::uuid)::text,$7)`, + event, + agent, + turn, + sequence, + input, + fresh, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + `UPDATE agent_inputs SET state='resolved',admitted_event_id=$2,admitted_at=$3,resolved_at=$3 WHERE id=$1`, + input, + event, + f.now, + ) + require.NoError(t, err) + if fresh { + _, err = tx.Exec( + ctx, + `INSERT INTO agent_turns(id,agent_id,turn_sequence,latest_event_id,latest_semantic_event_id) + VALUES($1,$2,$3,$4,$4)`, + turn, + agent, + sequence, + event, + ) + } else { + _, err = tx.Exec(ctx, ` +UPDATE agent_turns SET latest_event_id=$2,latest_semantic_event_id=$2 WHERE id=$1 +`, turn, event) + } + require.NoError(t, err) + require.NoError(t, tx.Commit(ctx)) + return turn, input +} + +func (f executionFixture) modelContext( + t *testing.T, + agent uuid.UUID, + watermark int64, + operation string, + source *int64, + attempt int, +) uuid.UUID { + t.Helper() + id := uuid.New() + _, err := f.pool.Exec( + t.Context(), + `INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,operation_kind,attempt_number, + agent_config_id,configured_model_revision_id,input_event_sequence,source_event_sequence_end,runtime_lock_id, +created_at) + VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)`, + id, + f.org, + f.project, + agent, + operation, + attempt, + f.config, + f.revision, + watermark, + source, + uuid.New(), + f.now, + ) + require.NoError(t, err) + return id +} + +func (f executionFixture) fail(t *testing.T, id uuid.UUID, recovery string) { + t.Helper() + var retry *time.Time + if recovery == "retry" { + retry = new(f.now.Add(time.Hour)) + } + _, err := f.pool.Exec( + t.Context(), + `UPDATE model_call_contexts SET state='failed',recovery_kind=$2,retry_at=$3, + error_kind='transient',completed_at=$4 WHERE id=$1`, + id, + recovery, + retry, + f.now, + ) + require.NoError(t, err) +} + +func (f executionFixture) output( + t *testing.T, + agent, turn, model uuid.UUID, + sequence int64, + reason string, + toolState string, +) (uuid.UUID, uuid.UUID) { + t.Helper() + ctx := t.Context() + tx, err := f.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + output, event, call := uuid.New(), uuid.New(), uuid.Nil + _, err = tx.Exec( + ctx, + `INSERT INTO model_outputs(id,agent_id,model_call_context_id,stop_reason,created_at) VALUES($1,$2,$3,$4,$5)`, + output, + agent, + model, + reason, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + `INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,model_output_id,created_at) + VALUES($1,$2,$3,$4,'model_output',$5,$6)`, + event, + agent, + turn, + sequence, + output, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec(ctx, `UPDATE agent_turns SET latest_event_id=$2, + latest_semantic_event_id=CASE WHEN $3='tool_use' THEN latest_semantic_event_id ELSE $2 END WHERE id=$1 +`, turn, event, reason) + require.NoError(t, err) + if toolState != "" { + call = uuid.New() + _, err = tx.Exec( + ctx, + `INSERT INTO tool_calls(id,agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) + VALUES($1,$2,$3,'call','list_agents','{}','built_in',$4,$5)`, + call, + agent, + output, + toolState, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + ` +INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,tool_call_id,created_at) + VALUES($1,'model_output',$2,0,'tool_call',$3,$4)`, + agent, + output, + call, + f.now, + ) + require.NoError(t, err) + } + _, err = tx.Exec( + ctx, + ` +UPDATE model_call_contexts SET state='succeeded',api_format='openai-responses',api_variant='default', +completed_at=$2 WHERE id=$1 +`, + model, + f.now, + ) + require.NoError(t, err) + require.NoError(t, tx.Commit(ctx)) + return output, call +} + +func (f executionFixture) complete(t *testing.T, agent, turn, call uuid.UUID, sequence int64) { + t.Helper() + ctx := t.Context() + tx, err := f.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + result, event := uuid.New(), uuid.New() + _, err = tx.Exec(ctx, `UPDATE tool_calls SET state='completed',runtime_lock_id=NULL WHERE id=$1`, call) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + `INSERT INTO tool_call_results(id,agent_id,tool_call_id,outcome,completed_at) VALUES($1,$2,$3,'succeeded',$4)`, + result, + agent, + call, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + `INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,tool_call_result_id,created_at) + VALUES($1,$2,$3,$4,'tool_result',$5,$6)`, + event, + agent, + turn, + sequence, + result, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec(ctx, `UPDATE agent_turns SET latest_event_id=$2 WHERE id=$1`, turn, event) + require.NoError(t, err) + require.NoError(t, tx.Commit(ctx)) +} + +func (f executionFixture) checkpoint( + t *testing.T, + agent, turn, model uuid.UUID, + sequence, source int64, +) uuid.UUID { + t.Helper() + ctx := t.Context() + tx, err := f.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + checkpoint, event := uuid.New(), uuid.New() + _, err = tx.Exec( + ctx, + ` +INSERT INTO context_checkpoints(id,agent_id,summarized_through_event_sequence,producer_model_call_context_id, +summary,created_at) + VALUES($1,$2,$3,$4,'summary',$5)`, + checkpoint, + agent, + source, + model, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + `INSERT INTO agent_events(id,agent_id,turn_id,sequence,event_kind,context_checkpoint_id,created_at) + VALUES($1,$2,$3,$4,'context_checkpoint',$5,$6)`, + event, + agent, + turn, + sequence, + checkpoint, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec( + ctx, + ` +UPDATE model_call_contexts SET state='succeeded',api_format='openai-responses',api_variant='default', +completed_at=$2 WHERE id=$1 +`, + model, + f.now, + ) + require.NoError(t, err) + _, err = tx.Exec(ctx, `UPDATE agent_turns SET latest_event_id=$2 WHERE id=$1`, turn, event) + require.NoError(t, err) + require.NoError(t, tx.Commit(ctx)) + return checkpoint +} + +func applyExecutionCutover(ctx context.Context, pool *pgxpool.Pool) error { + db := stdlib.OpenDBFromPool(pool) + defer func() { _ = db.Close() }() + return dbmigrate.ApplyPostgres(ctx, db, os.DirFS("../../../../migrations"), schemamigrations.GoMigrations()...) +} + +func (f executionFixture) handle( + t *testing.T, + agent uuid.UUID, +) (*agentexecution.Unit, *agentexecution.Handle) { + t.Helper() + ctx := t.Context() + cell := agentexecution.NewCell("test", f.pool, nil, nil) + u, err := cell.Begin(ctx) + require.NoError(t, err) + t.Cleanup(func() { _ = u.Rollback(context.Background()) }) + plan, err := u.PlanAgents( + ctx, + []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: agent}}, + agentexecution.RepairAuthority{}, + ) + require.NoError(t, err) + require.NoError(t, u.LockAgents(ctx, plan)) + var root uuid.UUID + require.NoError( + t, + u.DB().QueryRow(ctx, `SELECT root_agent_id FROM agents WHERE id=$1`, agent).Scan(&root), + ) + h, err := u.Agent( + agentexecution.AgentRoute{CellID: cell.ID(), ProjectID: f.project, AgentID: agent, RootAgentID: root}, + ) + require.NoError(t, err) + return u, h +} + +func TestReleasedExecutionReconstruction(t *testing.T) { + f := releasedExecutionFixture(t) + cases := []struct { + name string + build func(uuid.UUID) + work agentexecution.WorkKind + origin agentexecution.CandidateOrigin + }{ + {"empty", func(uuid.UUID) {}, agentexecution.WorkNone, 0}, + { + "initial", + func(a uuid.UUID) { f.opening(t, a, 1, uuid.Nil, "content") }, + agentexecution.WorkModel, + agentexecution.OriginInitial, + }, + {"retry", func(a uuid.UUID) { + f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + f.fail(t, c, "retry") + }, agentexecution.WorkNone, agentexecution.OriginRetry}, + { + "started", + func(a uuid.UUID) { + f.opening(t, a, 1, uuid.Nil, "content") + f.modelContext(t, a, 1, "normal", nil, 1) + }, + agentexecution.WorkNone, + 0, + }, + {"output_limit", func(a uuid.UUID) { + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + f.output(t, a, turn, c, 2, "max_tokens", "") + }, agentexecution.WorkModel, agentexecution.OriginOutputLimit}, + {"tools", func(a uuid.UUID) { + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + f.output(t, a, turn, c, 2, "tool_use", "awaiting_authorization") + }, agentexecution.WorkTool, 0}, + {"completed_tools", func(a uuid.UUID) { + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + _, call := f.output(t, a, turn, c, 2, "tool_use", "awaiting_authorization") + f.complete(t, a, turn, call, 3) + }, agentexecution.WorkModel, agentexecution.OriginTools}, + {"config", func(a uuid.UUID) { + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + f.output(t, a, turn, c, 2, "end_turn", "") + f.opening(t, a, 3, turn, "config_change") + }, agentexecution.WorkModel, agentexecution.OriginConfig}, + {"semantic", func(a uuid.UUID) { + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + f.fail(t, c, "retry") + f.opening(t, a, 2, turn, "config_change") + }, agentexecution.WorkModel, agentexecution.OriginSemantic}, + {"compaction", func(a uuid.UUID) { + f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + f.fail(t, c, "compact") + f.modelContext(t, a, 1, "compaction", new(int64(1)), 1) + }, agentexecution.WorkNone, 0}, + {"checkpoint", func(a uuid.UUID) { + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + normal := f.modelContext(t, a, 1, "normal", nil, 1) + f.fail(t, normal, "compact") + compact := f.modelContext(t, a, 1, "compaction", new(int64(1)), 1) + f.checkpoint(t, a, turn, compact, 2, 1) + }, agentexecution.WorkModel, agentexecution.OriginCheckpoint}, + } + + ids := make([]uuid.UUID, len(cases)) + for i, tc := range cases { + ids[i] = f.agent(t) + tc.build(ids[i]) + } + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + for i, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + u, h := f.handle(t, ids[i]) + got, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + rebuilt, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + got.Head.LogicalReadyAt = nil + rebuilt.Head.LogicalReadyAt = nil + require.Equal(t, got.Head, rebuilt.Head) + require.Equal(t, tc.work, got.Selection.Work) + + if tc.origin != 0 { + require.NotNil(t, got.Selection.Model) + require.Equal(t, tc.origin, got.Selection.Model.Origin) + } + require.NoError(t, u.Rollback(t.Context())) + }) + } +} + +func TestExecutionCutoverRejectsMultipleBatches(t *testing.T) { + f := releasedExecutionFixture(t) + agent := f.agent(t) + turn, _ := f.opening(t, agent, 1, uuid.Nil, "content") + first := f.modelContext(t, agent, 1, "normal", nil, 1) + f.output(t, agent, turn, first, 2, "tool_use", "awaiting_authorization") + second := f.modelContext(t, agent, 2, "normal", nil, 1) + f.output(t, agent, turn, second, 3, "tool_use", "awaiting_authorization") + err := applyExecutionCutover(t.Context(), f.pool) + require.ErrorContains(t, err, "multiple unconsumed tool outputs") + var exists bool + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT to_regclass('agent_execution_state') IS NOT NULL`).Scan(&exists), + ) + require.False(t, exists) +} + +func TestExecutionCutoverRejectsEmptyContextLineage(t *testing.T) { + f := releasedExecutionFixture(t) + agent := f.agent(t) + f.opening(t, agent, 1, uuid.Nil, "config_change") + f.modelContext(t, agent, 1, "normal", nil, 1) + require.Error(t, applyExecutionCutover(t.Context(), f.pool)) +} diff --git a/internal/storage/internal/agentexecution/repair.go b/internal/storage/internal/agentexecution/repair.go new file mode 100644 index 000000000..bd6048e85 --- /dev/null +++ b/internal/storage/internal/agentexecution/repair.go @@ -0,0 +1,216 @@ +package agentexecution + +import ( + "context" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" +) + +func (h *Handle) ReconstructExecution(ctx context.Context) (ExecutionSnapshot, error) { + l, err := h.executionLoader(ctx) + if err != nil { + return ExecutionSnapshot{}, err + } + valid, err := l.q.ValidateExecutionLineage( + ctx, + l.db, + executiondb.ValidateExecutionLineageParams{AgentID: l.route.AgentID}, + ) + if err != nil { + return ExecutionSnapshot{}, err + } + if !valid { + return ExecutionSnapshot{}, fmt.Errorf("%w: agent %s immutable lineage", ErrInvalidState, l.route.AgentID) + } + boundaries, err := l.q.LoadExecutionBoundaries(ctx, l.db, executiondb.LoadExecutionBoundariesParams{ + ProjectID: l.route.ProjectID, AgentID: l.route.AgentID}) + if err != nil { + return ExecutionSnapshot{}, err + } + head := ExecutionHead{ + AgentID: l.route.AgentID, + CurrentTurnID: boundaries.TurnID, + StopSequence: boundaries.StopSequence, + AnsweredThroughSequence: max(boundaries.StopSequence, boundaries.OutputSequence), + } + contexts, err := l.q.ReconstructExecutionContexts( + ctx, + l.db, + executiondb.ReconstructExecutionContextsParams{AgentID: head.AgentID}, + ) + if err != nil { + return ExecutionSnapshot{}, err + } + records := make(map[uuid.UUID]ModelContext, len(contexts)) + var normal, compaction *ModelContext + for _, row := range contexts { + record := contextRecord(row) + records[record.ID] = record + head.MaxContextInputSequence = max(head.MaxContextInputSequence, record.InputEventSequence) + if record.Operation != OperationNormal { + continue + } + head.MaxNormalInputSequence = max(head.MaxNormalInputSequence, record.InputEventSequence) + if record.TurnID == head.CurrentTurnID && + (normal == nil || record.InputEventSequence > normal.InputEventSequence || + record.InputEventSequence == normal.InputEventSequence && record.Attempt > normal.Attempt) { + normal = &record + } + } + if err := validateGoverningContexts(head, records); err != nil { + return ExecutionSnapshot{}, err + } + if normal != nil { + head.NormalContextID = normal.ID + for _, record := range records { + if record.TurnID != head.CurrentTurnID || record.Operation != OperationCompaction || + record.InputEventSequence != normal.InputEventSequence { + continue + } + if compaction == nil || record.SourceEventSequenceEnd < compaction.SourceEventSequenceEnd || + record.SourceEventSequenceEnd == compaction.SourceEventSequenceEnd && + record.Attempt > compaction.Attempt { + record := record + compaction = &record + } + } + } + if compaction != nil { + head.CompactionContextID = compaction.ID + } + if head.CurrentTurnID == uuid.Nil { + return l.load(ctx, head) + } + turn, err := l.turn(ctx, head) + if err != nil { + return ExecutionSnapshot{}, err + } + head.PendingConfigInputID, err = l.q.ReconstructExecutionConfig( + ctx, + l.db, + executiondb.ReconstructExecutionConfigParams{ + AgentID: head.AgentID, TurnID: head.CurrentTurnID, NormalWatermark: head.MaxNormalInputSequence, + StopSequence: head.StopSequence, FirstContentSequence: turn.FirstContentSequence}, + ) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return ExecutionSnapshot{}, err + } + head.PendingCheckpointID, err = l.q.ReconstructExecutionCheckpoint( + ctx, + l.db, + executiondb.ReconstructExecutionCheckpointParams{ + AgentID: head.AgentID, TurnID: head.CurrentTurnID, + NormalWatermark: head.MaxNormalInputSequence, StopSequence: head.StopSequence}, + ) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return ExecutionSnapshot{}, err + } + outputs, err := l.q.ReconstructExecutionOutputs(ctx, l.db, executiondb.ReconstructExecutionOutputsParams{ + AgentID: head.AgentID, TurnID: head.CurrentTurnID, StopSequence: head.StopSequence}) + if err != nil { + return ExecutionSnapshot{}, err + } + var acceptedWatermark int64 + for _, row := range outputs { + source, ok := records[row.ModelCallContextID] + if !ok || source.TurnID != head.CurrentTurnID || row.Sequence <= source.InputEventSequence { + return ExecutionSnapshot{}, fmt.Errorf("%w: output %s boundary", ErrInvalidState, row.ID) + } + acceptedWatermark = max(acceptedWatermark, source.InputEventSequence) + } + for _, row := range outputs { + source, ok := records[row.ModelCallContextID] + if !ok { + return ExecutionSnapshot{}, fmt.Errorf("%w: output %s missing context", ErrInvalidState, row.ID) + } + output := ModelOutput{ + ID: row.ID, + EventID: row.EventID, + Event: EventBoundary{Sequence: row.Sequence, Time: row.CreatedAt}, + Context: source, + } + batch, loadErr := l.batch(ctx, output) + if loadErr != nil { + return ExecutionSnapshot{}, loadErr + } + if batch != nil { + if batch.Completion != nil && acceptedWatermark >= batch.Completion.LastResultSequence { + continue + } + if head.PendingToolOutputID != uuid.Nil { + return ExecutionSnapshot{}, fmt.Errorf( + "%w: multiple tool outputs %s and %s", + ErrInvalidState, + head.PendingToolOutputID, + row.ID, + ) + } + head.PendingToolOutputID = row.ID + } else if row.StopReason == "max_tokens" && source.Operation == OperationNormal && source.State == ContextSucceeded && + row.Sequence > head.MaxNormalInputSequence { + if head.PendingOutputLimitID != uuid.Nil { + return ExecutionSnapshot{}, fmt.Errorf("%w: multiple output limits %s and %s", + ErrInvalidState, head.PendingOutputLimitID, row.ID) + } + head.PendingOutputLimitID = row.ID + } + } + return l.load(ctx, head) +} + +func validateGoverningContexts(head ExecutionHead, records map[uuid.UUID]ModelContext) error { + type identity struct { + operation Operation + watermark int64 + end int64 + } + attempts := make(map[identity]int32) + ends := make(map[int64]int64) + var normalWatermark int64 + for _, record := range records { + if record.TurnID != head.CurrentTurnID { + continue + } + key := identity{record.Operation, record.InputEventSequence, record.SourceEventSequenceEnd} + attempts[key] = max(attempts[key], record.Attempt) + if record.Operation == OperationNormal { + normalWatermark = max(normalWatermark, record.InputEventSequence) + } else { + end, ok := ends[record.InputEventSequence] + if !ok || record.SourceEventSequenceEnd < end { + ends[record.InputEventSequence] = record.SourceEventSequenceEnd + } + } + } + var governing uuid.UUID + for _, record := range records { + if record.TurnID != head.CurrentTurnID || record.InputEventSequence < normalWatermark || + record.InputEventSequence < head.StopSequence { + continue + } + key := identity{record.Operation, record.InputEventSequence, record.SourceEventSequenceEnd} + if record.Attempt != attempts[key] { + continue + } + end, dependency := ends[record.InputEventSequence] + if record.Operation == OperationNormal && dependency || + record.Operation == OperationCompaction && record.SourceEventSequenceEnd != end { + continue + } + if record.State != ContextStarted && (record.State != ContextFailed || record.Recovery == RecoveryNone) { + continue + } + if governing != uuid.Nil { + return fmt.Errorf("%w: governing contexts %s and %s", ErrInvalidState, governing, record.ID) + } + governing = record.ID + if record.Operation == OperationCompaction && record.InputEventSequence != normalWatermark { + return fmt.Errorf("%w: compaction %s without normal parent", ErrInvalidState, record.ID) + } + } + return nil +} diff --git a/internal/storage/internal/agentexecution/route.go b/internal/storage/internal/agentexecution/route.go new file mode 100644 index 000000000..234002b37 --- /dev/null +++ b/internal/storage/internal/agentexecution/route.go @@ -0,0 +1,23 @@ +package agentexecution + +import ( + "bytes" + + "github.com/google/uuid" +) + +type CellID string + +type AgentRoute struct { + CellID CellID + ProjectID uuid.UUID + AgentID uuid.UUID + RootAgentID uuid.UUID +} + +func compareRoutes(a, b AgentRoute) int { + if a.ProjectID != b.ProjectID { + return bytes.Compare(a.ProjectID[:], b.ProjectID[:]) + } + return bytes.Compare(a.AgentID[:], b.AgentID[:]) +} diff --git a/internal/storage/internal/agentexecution/schema_integration_test.go b/internal/storage/internal/agentexecution/schema_integration_test.go new file mode 100644 index 000000000..456a00637 --- /dev/null +++ b/internal/storage/internal/agentexecution/schema_integration_test.go @@ -0,0 +1,317 @@ +//go:build integration + +package agentexecution_test + +import ( + "errors" + "go/ast" + "go/parser" + "go/token" + "os" + "strconv" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func TestExecutionSchemaContracts(t *testing.T) { + f := releasedExecutionFixture(t) + a, b := f.agent(t), f.agent(t) + turn, _ := f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + output, _ := f.output(t, a, turn, c, 2, "max_tokens", "") + otherTurn, _ := f.opening(t, b, 1, uuid.Nil, "content") + otherContext := f.modelContext(t, b, 1, "normal", nil, 1) + otherOutput, _ := f.output(t, b, otherTurn, otherContext, 2, "max_tokens", "") + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + var columns int + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM information_schema.columns + WHERE table_schema='public' AND table_name='agent_execution_state' + AND (data_type IN ('json','jsonb') OR column_name IN ('project_id','dirty','incremental'))`).Scan(&columns)) + require.Zero(t, columns) + var fks int + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM pg_constraint + WHERE conrelid='agent_execution_state'::regclass AND contype='f'`).Scan(&fks)) + require.Equal(t, 8, fks) + for _, index := range []string{"agents_root_idx", + "tool_calls_output_idx", + "tool_calls_output_incomplete_idx", + "tool_calls_output_runnable_idx", "model_call_contexts_turn_kind_idx", "agent_events_output_boundary_idx"} { + var valid bool + require.NoError( + t, + f.pool.QueryRow(t.Context(), ` +SELECT indisvalid AND indisready FROM pg_index WHERE indexrelid=$1::regclass +`, index). + Scan(&valid), + ) + require.True(t, valid, index) + } + for _, tc := range []struct { + column string + value uuid.UUID + }{ + {"current_turn_id", + otherTurn}, + {"normal_context_id", + otherContext}, + {"pending_output_limit_id", + otherOutput}, + } { + t.Run(tc.column, func(t *testing.T) { + _, err := f.pool.Exec( + t.Context(), + `UPDATE agent_execution_state SET `+tc.column+`=$2 WHERE agent_id=$1`, + a, + tc.value, + ) + var constraint *pgconn.PgError + require.ErrorAs(t, err, &constraint) + require.Equal(t, "23503", constraint.Code) + }) + } + _, err := f.pool.Exec(t.Context(), `UPDATE model_call_contexts SET opening_input_ids='{}' WHERE id=$1`, c) + require.Error(t, err) + _, err = f.pool.Exec(t.Context(), `UPDATE agents SET root_agent_id=$2 WHERE id=$1`, a, b) + require.Error(t, err) + _, err = f.pool.Exec( + t.Context(), + ` +INSERT INTO agents(id,root_agent_id,parent_agent_id,subagent_key,org_id,project_id,state,current_config_id, +created_at,updated_at) + VALUES($1,$2,$3,'child',$4,$5,'active',$6,$7,$7)`, + uuid.New(), + b, + a, + f.org, + f.project, + f.config, + f.now, + ) + require.ErrorContains(t, err, "root must match parent") + child := uuid.New() + _, err = f.pool.Exec( + t.Context(), + ` +INSERT INTO agents(id,root_agent_id,parent_agent_id,subagent_key,org_id,project_id,state,current_config_id, +created_at,updated_at) + VALUES($1,$2,$2,'child',$3,$4,'active',$5,$6,$6)`, + child, + a, + f.org, + f.project, + f.config, + f.now, + ) + require.NoError(t, err) + var childRoot uuid.UUID + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT root_agent_id FROM agents WHERE id=$1`, child).Scan(&childRoot), + ) + require.Equal(t, a, childRoot) + require.NotEqual(t, output, otherOutput) + for _, tc := range []struct { + ids []uuid.UUID + sequence *int64 + valid bool + }{ + {[]uuid.UUID{}, nil, true}, {[]uuid.UUID{uuid.New()}, new(int64(1)), true}, + {nil, nil, false}, {[]uuid.UUID{}, new(int64(1)), false}, {[]uuid.UUID{uuid.New()}, nil, false}, + {[]uuid.UUID{a, a}, new(int64(1)), false}, {[]uuid.UUID{uuid.Nil}, new(int64(1)), false}, + } { + var valid bool + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT execution_opening_is_valid($1,$2)`, tc.ids, tc.sequence). + Scan(&valid), + ) + require.Equal(t, tc.valid, valid) + } +} + +func TestExecutionReconstructionIgnoresCorruptHead(t *testing.T) { + f := releasedExecutionFixture(t) + agent := f.agent(t) + turn, input := f.opening(t, agent, 1, uuid.Nil, "content") + c := f.modelContext(t, agent, 1, "normal", nil, 1) + output, _ := f.output(t, agent, turn, c, 2, "max_tokens", "") + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + u, h := f.handle(t, agent) + expected, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + _, err = u.DB(). + Exec(t.Context(), `UPDATE agent_execution_state SET normal_context_id=NULL,pending_output_limit_id=NULL, + max_normal_input_sequence=0,max_context_input_sequence=0,turn_continuable=false,logical_ready_at=NULL WHERE +agent_id=$1 +`, agent) + require.NoError(t, err) + for range 2 { + rebuilt, rebuildErr := h.ReconstructExecution(t.Context()) + require.NoError(t, rebuildErr) + require.Equal(t, expected.Head, rebuilt.Head) + require.Equal(t, output, rebuilt.Head.PendingOutputLimitID) + require.Equal(t, []uuid.UUID{input}, rebuilt.Selection.Model.Opening.InputIDs) + } + var unchanged bool + require.NoError( + t, + u.DB(). + QueryRow(t.Context(), ` +SELECT normal_context_id IS NULL FROM agent_execution_state WHERE agent_id=$1 +`, agent). + Scan(&unchanged), + ) + require.True(t, unchanged) + rejected := errors.New("rollback input") + require.ErrorIs(t, u.Savepoint(t.Context(), func(u *agentexecution.Unit) error { + _, insertErr := u.DB(). + Exec(t.Context(), `INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) + VALUES($1,$2,'received','content','steering',$3)`, f.project, agent, f.now) + require.NoError(t, insertErr) + rebuilt, rebuildErr := h.ReconstructExecution(t.Context()) + require.NoError(t, rebuildErr) + require.Equal(t, agentexecution.AdmitAllSteering, rebuilt.Selection.Admission) + return rejected + }), rejected) + rebuilt, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + require.Equal(t, agentexecution.WorkModel, rebuilt.Selection.Work) + _, err = u.DB().Exec(t.Context(), `DELETE FROM agent_execution_state WHERE agent_id=$1`, agent) + require.NoError(t, err) + _, err = h.LoadExecution(t.Context()) + require.ErrorIs(t, err, pgx.ErrNoRows) + rebuilt, err = h.ReconstructExecution(t.Context()) + require.NoError(t, err) + require.Equal(t, expected.Head, rebuilt.Head) +} + +func TestExecutionReconstructionRejectsCorruptLineage(t *testing.T) { + f := releasedExecutionFixture(t) + a, b := f.agent(t), f.agent(t) + f.opening(t, a, 1, uuid.Nil, "content") + c := f.modelContext(t, a, 1, "normal", nil, 1) + _, other := f.opening(t, b, 1, uuid.Nil, "content") + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + u, h := f.handle(t, a) + _, err := u.DB().Exec(t.Context(), `ALTER TABLE model_call_contexts DISABLE TRIGGER USER`) + require.NoError(t, err) + _, err = u.DB(). + Exec(t.Context(), `UPDATE model_call_contexts SET opening_input_ids=ARRAY[$2::uuid] WHERE id=$1`, c, other) + require.NoError(t, err) + _, err = h.ReconstructExecution(t.Context()) + require.ErrorIs(t, err, agentexecution.ErrInvalidState) +} + +func TestExecutionFactPlans(t *testing.T) { + f := releasedExecutionFixture(t) + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + entries, err := os.ReadDir("internal/executiondb") + require.NoError(t, err) + connection, err := f.pool.Acquire(t.Context()) + require.NoError(t, err) + defer connection.Release() + for _, entry := range entries { + if !strings.HasSuffix(entry.Name(), ".sql.go") && entry.Name() != "batch.go" { + continue + } + file, parseErr := parser.ParseFile(token.NewFileSet(), "internal/executiondb/"+entry.Name(), nil, 0) + require.NoError(t, parseErr) + ast.Inspect(file, func(node ast.Node) bool { + literal, ok := node.(*ast.BasicLit) + if !ok || literal.Kind != token.STRING { + return true + } + query, unquoteErr := strconv.Unquote(literal.Value) + require.NoError(t, unquoteErr) + if !strings.HasPrefix(query, "-- name: Load") && !strings.HasPrefix(query, "-- name: Capture") && + !strings.HasPrefix( + query, + "-- name: PublishExecution ", + ) && + !strings.HasPrefix(query, + "-- name: ListExecutionIntegrationTargetAgents ") && + !strings.HasPrefix(query, + "-- name: ClearExecutionIntegrationTargets ") { + return true + } + name := strings.Fields(query)[2] + t.Run(name, func(t *testing.T) { + results, explainErr := connection.Conn(). + PgConn(). + Exec(t.Context(), "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+query). + ReadAll() + require.NoError(t, explainErr) + require.Len(t, results, 1) + plan := string(results[0].Rows[0][0]) + require.NotContains(t, plan, "Function Scan") + limit := 22 + if name == "LoadExecutionFacts" { + limit = 100 + } + if name == "PublishExecution" { + limit = 35 + } + if name == "LoadExecutionBase" { + limit = 28 + } + require.LessOrEqual(t, strings.Count(plan, `"Node Type"`), limit) + require.NotContains(t, plan, "agent_next_") + require.NotContains(t, plan, "scheduler_") + if name == "LoadExecutionBase" || name == "LoadExecutionFacts" || + name == "ListExecutionIntegrationTargetAgents" || + name == "ClearExecutionIntegrationTargets" { + _, err := connection.Exec(t.Context(), "SET enable_seqscan=off") + require.NoError(t, err) + indexed, err := connection.Conn().PgConn(). + Exec(t.Context(), "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+query).ReadAll() + require.NoError(t, err) + require.Len(t, indexed, 1) + require.NotContains(t, string(indexed[0].Rows[0][0]), `"Node Type": "Seq Scan"`) + _, err = connection.Exec(t.Context(), "RESET enable_seqscan") + require.NoError(t, err) + t.Logf("generic plan: %d nodes", strings.Count(plan, `"Node Type"`)) + } + }) + return true + }) + } +} + +func TestExecutionRootBackfill(t *testing.T) { + f := releasedExecutionFixture(t) + root := f.agent(t) + child, grandchild := uuid.New(), uuid.New() + _, err := f.pool.Exec(t.Context(), `INSERT INTO agents(id,parent_agent_id,subagent_key, + org_id,project_id,state,current_config_id,created_at,updated_at) + VALUES($1,$3,'child',$4,$5,'active',$6,$7,$7),($2,$1,'grandchild',$4,$5,'active',$6,$7,$7)`, + child, grandchild, root, f.org, f.project, f.config, f.now) + require.NoError(t, err) + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + var roots []uuid.UUID + require.NoError( + t, + f.pool.QueryRow(t.Context(), `SELECT array_agg(root_agent_id) FROM agents`).Scan(&roots), + ) + require.Equal(t, []uuid.UUID{root, root, root}, roots) + _, h := f.handle(t, grandchild) + snapshot, err := h.ReconstructExecution(t.Context()) + require.NoError(t, err) + require.Equal(t, grandchild, snapshot.Head.AgentID) +} + +func TestExecutionRootBackfillRejectsCycle(t *testing.T) { + f := releasedExecutionFixture(t) + a, b := uuid.New(), uuid.New() + _, err := f.pool.Exec(t.Context(), `INSERT INTO agents(id,parent_agent_id,subagent_key, + org_id,project_id,state,current_config_id,created_at,updated_at) + VALUES($1,$2,'a',$3,$4,'active',$5,$6,$6),($2,$1,'b',$3,$4,'active',$5,$6,$6)`, + a, b, f.org, f.project, f.config, f.now) + require.NoError(t, err) + require.ErrorContains(t, applyExecutionCutover(t.Context(), f.pool), "invalid agent ancestry") +} diff --git a/internal/storage/internal/agentexecution/selection.go b/internal/storage/internal/agentexecution/selection.go new file mode 100644 index 000000000..11a697f1d --- /dev/null +++ b/internal/storage/internal/agentexecution/selection.go @@ -0,0 +1,283 @@ +package agentexecution + +import ( + "errors" + "fmt" + "slices" + "time" + + "github.com/google/uuid" +) + +var ErrInvalidState = errors.New("invalid agent execution state") + +func SelectNext(view ExecutionView, now time.Time) (Selection, error) { + if err := validateView(view); err != nil { + return Selection{}, err + } + if view.State == AgentArchived { + return Selection{Wait: WaitArchived}, nil + } + selection := Selection{Wait: WaitIdle} + governing := continuableContext(view) + selection.TurnContinuable = governing != nil + if batch := view.ToolBatch; batch != nil && view.StopSequence <= batch.Output.Event.Sequence { + selection.IncompleteTools = batch.HasIncomplete + if batch.HasIncomplete { + selection.TurnContinuable = true + selection.Wait = WaitToolBatch + if batch.HasRunnable { + selection.Work = WorkTool + selection.Wait = WaitNone + selection.LogicalReadyAt = &now + selection.Tool = &ToolDecision{ + TurnID: batch.Output.Context.TurnID, SourceContextID: batch.Output.Context.ID, + OutputID: batch.Output.ID, SourceEventID: batch.Output.EventID, + } + } + return selection, nil + } + } + selection.Model = selectModel(view, governing) + selection.TurnContinuable = selection.TurnContinuable || selection.Model != nil + switch { + case view.Inputs.Steering: + selection.Work, selection.Admission = WorkInput, AdmitAllSteering + selection.LogicalReadyAt = &now + case selection.Model != nil: + selection.LogicalReadyAt = &selection.Model.ReadyAt + if selection.Model.ReadyAt.After(now) { + selection.Wait = WaitModelDeadline + return selection, nil + } + selection.Work = WorkModel + case view.Inputs.Queued: + selection.Work, selection.Admission = WorkInput, AdmitOneQueued + selection.LogicalReadyAt = &now + default: + return selection, nil + } + selection.Wait = WaitNone + return selection, nil +} + +func continuableContext(view ExecutionView) *ModelContext { + governing := view.NormalContext + if view.CompactionContext != nil { + governing = view.CompactionContext + } + if governing == nil || view.StopSequence > governing.InputEventSequence { + return nil + } + if governing.State == ContextStarted || (governing.State == ContextFailed && governing.Recovery != RecoveryNone) { + return governing + } + return nil +} + +func selectModel(view ExecutionView, governing *ModelContext) *ModelDecision { + turn := view.Turn + if turn == nil { + return nil + } + candidates := make([]ModelDecision, 0, 7) + if view.MaxContextInputSequence < turn.FirstOpeningSequence && view.StopSequence <= turn.FirstOpeningSequence { + candidates = append(candidates, ModelDecision{ + Kind: ModelStart, Origin: OriginInitial, TurnID: turn.ID, + Opening: turn.InitialOpening, ReadyAt: turn.InitialReadyAt, + }) + } + if governing != nil { + laterSemantic := turn.LatestSemantic.Sequence > governing.InputEventSequence + if !laterSemantic && governing.State == ContextFailed && governing.Recovery == RecoveryRetry { + candidates = append(candidates, ModelDecision{ + Kind: ModelResume, Origin: OriginRetry, TurnID: turn.ID, SourceContextID: governing.ID, + Opening: governing.Opening, ReadyAt: *governing.RetryAt, + }) + } + if laterSemantic { + candidates = append(candidates, ModelDecision{ + Kind: ModelStart, Origin: OriginSemantic, TurnID: turn.ID, SourceContextID: governing.ID, + Opening: governing.Opening, ReadyAt: turn.LatestSemantic.Time, + }) + } + } + if config := view.Config; eligibleBoundary(view, config) && + turn.FirstContentSequence > 0 && config.Event.Sequence >= turn.FirstContentSequence { + candidates = append(candidates, ModelDecision{ + Kind: ModelStart, Origin: OriginConfig, TurnID: turn.ID, SourceInputID: config.ID, + Opening: config.Opening, ReadyAt: config.Event.Time, + }) + } + if batch := view.ToolBatch; batch != nil && view.StopSequence <= batch.Output.Event.Sequence { + candidates = append(candidates, ModelDecision{ + Kind: ModelContinue, Origin: OriginTools, TurnID: turn.ID, + SourceContextID: batch.Output.Context.ID, SourceOutputID: batch.Output.ID, + Opening: batch.Output.Context.Opening, ReadyAt: batch.Completion.ReadyAt, + }) + } + if output := view.OutputLimit; output != nil && view.StopSequence <= output.Event.Sequence && + output.Event.Sequence > view.MaxNormalInputSequence { + candidates = append(candidates, ModelDecision{ + Kind: ModelContinue, Origin: OriginOutputLimit, TurnID: turn.ID, + SourceContextID: output.Context.ID, SourceOutputID: output.ID, + Opening: output.Context.Opening, ReadyAt: output.Event.Time, + }) + } + if checkpoint := view.Checkpoint; eligibleBoundary(view, checkpoint) { + candidates = append(candidates, ModelDecision{ + Kind: ModelStart, Origin: OriginCheckpoint, TurnID: turn.ID, SourceCheckpointID: checkpoint.ID, + Opening: checkpoint.Opening, ReadyAt: checkpoint.Event.Time, + }) + } + for _, candidate := range candidates { + if len(candidate.Opening.InputIDs) != 0 { + candidate.Opening.InputIDs = slices.Clone(candidate.Opening.InputIDs) + return &candidate + } + } + return nil +} + +func eligibleBoundary(view ExecutionView, work *BoundaryWork) bool { + return work != nil && work.Event.Sequence > view.MaxNormalInputSequence && + view.StopSequence <= work.Event.Sequence +} + +func validateView(view ExecutionView) error { + invalid := func(reason string) error { return fmt.Errorf("%w: %s", ErrInvalidState, reason) } + if view.AgentID == uuid.Nil || (view.State != AgentActive && view.State != AgentArchived) { + return invalid("agent identity or lifecycle") + } + if view.StopSequence < 0 || view.MaxNormalInputSequence < 0 || + view.MaxContextInputSequence < view.MaxNormalInputSequence { + return invalid("event watermarks") + } + if view.Turn == nil { + if view.NormalContext != nil || view.CompactionContext != nil || view.ToolBatch != nil || + view.OutputLimit != nil || view.Config != nil || view.Checkpoint != nil { + return invalid("work without a current turn") + } + return nil + } + turn := view.Turn + if turn.ID == uuid.Nil || turn.FirstOpeningSequence <= 0 || + turn.LastOpeningSequence < turn.FirstOpeningSequence { + return invalid("turn opening boundary") + } + if err := validateOpening(turn.InitialOpening); err != nil { + return err + } + for _, governing := range []*ModelContext{view.NormalContext, view.CompactionContext} { + if governing == nil { + continue + } + if err := validateContext(view, *governing); err != nil { + return err + } + } + if governing := view.NormalContext; governing != nil && + (governing.Operation != OperationNormal || governing.InputEventSequence != view.MaxNormalInputSequence) { + return invalid("normal context head") + } + if governing := view.CompactionContext; governing != nil && + (governing.Operation != OperationCompaction || view.NormalContext == nil || + governing.InputEventSequence != view.NormalContext.InputEventSequence) { + return invalid("compaction dependency head") + } + if view.ToolBatch != nil && view.OutputLimit != nil { + return invalid("multiple pending model outputs") + } + if batch := view.ToolBatch; batch != nil { + if err := validateOutput(view, batch.Output); err != nil { + return err + } + if (batch.HasRunnable && !batch.HasIncomplete) || (batch.HasIncomplete == (batch.Completion != nil)) { + return invalid("tool batch completion facts") + } + if batch.Completion != nil && batch.Completion.LastResultSequence <= batch.Output.Event.Sequence { + return invalid("tool result frontier") + } + } + if output := view.OutputLimit; output != nil { + if err := validateOutput(view, *output); err != nil { + return err + } + } + for _, work := range []*BoundaryWork{view.Config, view.Checkpoint} { + if work == nil { + continue + } + if work.ID == uuid.Nil || work.AgentID != view.AgentID || work.TurnID != turn.ID || work.Event.Sequence <= 0 { + return invalid("config or checkpoint scope") + } + if err := validateOpening(work.Opening); err != nil { + return err + } + } + return nil +} + +func validateContext(view ExecutionView, governing ModelContext) error { + if governing.ID == uuid.Nil || governing.AgentID != view.AgentID || governing.TurnID != view.Turn.ID || + governing.Attempt < 1 || governing.InputEventSequence < view.Turn.FirstOpeningSequence || + governing.InputEventSequence > view.MaxContextInputSequence { + return fmt.Errorf("%w: context scope or identity", ErrInvalidState) + } + if (governing.Operation == OperationNormal && governing.SourceEventSequenceEnd != 0) || + (governing.Operation == OperationCompaction && (governing.SourceEventSequenceEnd <= 0 || + governing.SourceEventSequenceEnd > governing.InputEventSequence)) || + (governing.Operation != OperationNormal && governing.Operation != OperationCompaction) { + return fmt.Errorf("%w: context operation", ErrInvalidState) + } + switch governing.State { + case ContextStarted, ContextSucceeded, ContextFailed, ContextCanceled: + default: + return fmt.Errorf("%w: context state", ErrInvalidState) + } + switch governing.Recovery { + case RecoveryNone, RecoveryRetry: + case RecoveryCompact: + if governing.Operation != OperationNormal { + return fmt.Errorf("%w: normal compaction dependency", ErrInvalidState) + } + case RecoveryReduce: + if governing.Operation != OperationCompaction { + return fmt.Errorf("%w: compaction source reduction", ErrInvalidState) + } + default: + return fmt.Errorf("%w: context recovery", ErrInvalidState) + } + if governing.Opening.EventSequence > governing.InputEventSequence { + return fmt.Errorf("%w: opening after context watermark", ErrInvalidState) + } + if governing.Recovery != RecoveryNone && governing.State != ContextFailed { + return fmt.Errorf("%w: recovery on nonfailed context", ErrInvalidState) + } + if (governing.Recovery == RecoveryRetry) != (governing.RetryAt != nil) { + return fmt.Errorf("%w: retry deadline", ErrInvalidState) + } + return validateOpening(governing.Opening) +} + +func validateOutput(view ExecutionView, output ModelOutput) error { + if output.ID == uuid.Nil || output.EventID == uuid.Nil || output.Context.Operation != OperationNormal || + output.Context.State != ContextSucceeded || output.Event.Sequence <= output.Context.InputEventSequence { + return fmt.Errorf("%w: pending model output", ErrInvalidState) + } + return validateContext(view, output.Context) +} + +func validateOpening(opening Opening) error { + if (len(opening.InputIDs) == 0) != (opening.EventSequence == 0) || opening.EventSequence < 0 { + return fmt.Errorf("%w: opening boundary", ErrInvalidState) + } + seen := make(map[uuid.UUID]struct{}, len(opening.InputIDs)) + for _, id := range opening.InputIDs { + if _, duplicate := seen[id]; duplicate || id == uuid.Nil { + return fmt.Errorf("%w: duplicate or empty opening identity", ErrInvalidState) + } + seen[id] = struct{}{} + } + return nil +} diff --git a/internal/storage/internal/agentexecution/selection_invariants_test.go b/internal/storage/internal/agentexecution/selection_invariants_test.go new file mode 100644 index 000000000..41fd82de9 --- /dev/null +++ b/internal/storage/internal/agentexecution/selection_invariants_test.go @@ -0,0 +1,191 @@ +package agentexecution_test + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func TestSelectNextRejectsMalformedViews(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name string + change func(*agentexecution.ExecutionView) + }{ + {"missing agent", func(v *agentexecution.ExecutionView) { v.AgentID = uuid.Nil }}, + {"unknown lifecycle", func(v *agentexecution.ExecutionView) { v.State = "unknown" }}, + {"negative stop", func(v *agentexecution.ExecutionView) { v.StopSequence = -1 }}, + {"normal beyond all contexts", func(v *agentexecution.ExecutionView) { v.MaxNormalInputSequence++ }}, + {"missing turn", func(v *agentexecution.ExecutionView) { v.Turn = nil }}, + {"missing turn id", func(v *agentexecution.ExecutionView) { v.Turn.ID = uuid.Nil }}, + {"reversed openings", func(v *agentexecution.ExecutionView) { v.Turn.LastOpeningSequence-- }}, + {"wrong agent context", func(v *agentexecution.ExecutionView) { v.NormalContext.AgentID = identity(99) }}, + {"wrong turn context", func(v *agentexecution.ExecutionView) { v.NormalContext.TurnID = identity(99) }}, + {"context before turn", func(v *agentexecution.ExecutionView) { v.NormalContext.InputEventSequence-- }}, + {"bad attempt", func(v *agentexecution.ExecutionView) { v.NormalContext.Attempt = 0 }}, + {"unknown context state", func(v *agentexecution.ExecutionView) { v.NormalContext.State = "unknown" }}, + {"unknown recovery", func(v *agentexecution.ExecutionView) { v.NormalContext.Recovery = "unknown" }}, + {"recovery on success", func(v *agentexecution.ExecutionView) { + v.NormalContext.Recovery = agentexecution.RecoveryCompact + }}, + {"retry without deadline", func(v *agentexecution.ExecutionView) { + v.NormalContext.State, v.NormalContext.Recovery = agentexecution.ContextFailed, agentexecution.RecoveryRetry + }}, + {"deadline without retry", func(v *agentexecution.ExecutionView) { + deadline := decisionTime() + v.NormalContext.RetryAt = &deadline + }}, + {"normal with range", func(v *agentexecution.ExecutionView) { v.NormalContext.SourceEventSequenceEnd = 1 }}, + {"duplicate opening", func(v *agentexecution.ExecutionView) { + v.NormalContext.Opening.InputIDs = []uuid.UUID{identity(3), identity(3)} + }}, + {"nil opening identity", func(v *agentexecution.ExecutionView) { + v.NormalContext.Opening.InputIDs = []uuid.UUID{uuid.Nil} + }}, + {"opening past watermark", func(v *agentexecution.ExecutionView) { v.NormalContext.Opening.EventSequence++ }}, + {"missing opening boundary", func(v *agentexecution.ExecutionView) { v.NormalContext.Opening.EventSequence = 0 }}, + {"orphan compaction", func(v *agentexecution.ExecutionView) { + *v = compactionView() + v.NormalContext = nil + }}, + {"compaction wrong watermark", func(v *agentexecution.ExecutionView) { + *v = compactionView() + v.CompactionContext.InputEventSequence++ + v.MaxContextInputSequence++ + }}, + {"wrong config scope", func(v *agentexecution.ExecutionView) { + v.Config = boundary(*v, 8, 20) + v.Config.AgentID = identity(99) + }}, + {"runnable completed batch", func(v *agentexecution.ExecutionView) { + *v = completedBatchView() + v.ToolBatch.HasRunnable = true + }}, + {"incomplete with completion", func(v *agentexecution.ExecutionView) { + *v = completedBatchView() + v.ToolBatch.HasIncomplete = true + }}, + {"complete without frontier", func(v *agentexecution.ExecutionView) { + *v = completedBatchView() + v.ToolBatch.Completion = nil + }}, + {"result before source", func(v *agentexecution.ExecutionView) { + *v = completedBatchView() + v.ToolBatch.Completion.LastResultSequence = v.ToolBatch.Output.Event.Sequence + }}, + {"two pending outputs", func(v *agentexecution.ExecutionView) { + *v = completedBatchView() + o := output(*v) + v.OutputLimit = &o + }}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + view := baseView() + test.change(&view) + _, err := agentexecution.SelectNext(view, decisionTime()) + require.ErrorIs(t, err, agentexecution.ErrInvalidState) + }) + } +} + +func TestCorrectedOpeningLineageSurvivesRepeatedContinuationAndRetry(t *testing.T) { + t.Parallel() + for _, tools := range []bool{false, true} { + view := completedBatchView() + inherited := agentexecution.Opening{InputIDs: []uuid.UUID{identity(11), identity(3)}, EventSequence: 2} + view.ToolBatch.Output.Context.Opening = inherited + if !tools { + o := view.ToolBatch.Output + view.ToolBatch, view.OutputLimit = nil, &o + } + for cycle := range 3 { + selection, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, agentexecution.ModelContinue, selection.Model.Kind) + require.Equal(t, inherited, selection.Model.Opening) + next := *view.NormalContext + next.ID, next.Attempt = identity(byte(20+cycle)), 1 + next.InputEventSequence = 30 + int64(cycle)*20 + next.Opening = selection.Model.Opening + next.State, next.Recovery = agentexecution.ContextFailed, agentexecution.RecoveryRetry + deadline := decisionTime() + next.RetryAt = &deadline + view.NormalContext, view.OutputLimit = &next, nil + view.MaxNormalInputSequence, view.MaxContextInputSequence = next.InputEventSequence, next.InputEventSequence + selection, err = agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, agentexecution.OriginRetry, selection.Model.Origin) + require.Equal(t, inherited, selection.Model.Opening) + next.ID, next.Attempt = identity(byte(40+cycle)), 2 + next.State, next.Recovery, next.RetryAt = agentexecution.ContextSucceeded, agentexecution.RecoveryNone, nil + o := agentexecution.ModelOutput{ + ID: identity(byte(60 + cycle)), EventID: identity(byte(80 + cycle)), Context: next, + Event: agentexecution.EventBoundary{Sequence: next.InputEventSequence + 1, Time: decisionTime()}, + } + if tools { + view.ToolBatch = &agentexecution.ToolBatch{Output: o, Completion: &agentexecution.BatchCompletion{ + LastResultSequence: o.Event.Sequence + 1, ReadyAt: decisionTime(), + }} + } else { + view.OutputLimit = &o + } + } + } +} + +func FuzzSelectNext(f *testing.F) { + for mode := range byte(6) { + for flags := range byte(16) { + f.Add(mode, flags) + } + } + f.Fuzz(func(t *testing.T, mode, flags byte) { + makeView := func() agentexecution.ExecutionView { + var view agentexecution.ExecutionView + switch mode % 6 { + case 0: + view = initialView() + case 1: + view = retryView() + case 2: + view = completedBatchView() + case 3: + view = batchView() + case 4: + view = batchView() + view.ToolBatch.HasRunnable = false + case 5: + view = compactionView() + } + view.Inputs.Steering, view.Inputs.Queued = flags&1 != 0, flags&2 != 0 + if flags&4 != 0 { + view.StopSequence = 40 + } + if flags&8 != 0 { + view.State = agentexecution.AgentArchived + } + return view + } + view, original := makeView(), makeView() + got, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, original, view) + again, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, got, again) + if got.IncompleteTools { + require.Nil(t, got.Model) + require.Equal(t, agentexecution.AdmitNone, got.Admission) + } + if got.Work == agentexecution.WorkModel { + require.False(t, got.Model.ReadyAt.After(decisionTime())) + later, err := agentexecution.SelectNext(view, decisionTime().Add(time.Second)) + require.NoError(t, err) + require.Equal(t, got.Model, later.Model) + } + }) +} diff --git a/internal/storage/internal/agentexecution/selection_load.go b/internal/storage/internal/agentexecution/selection_load.go new file mode 100644 index 000000000..a05a6482d --- /dev/null +++ b/internal/storage/internal/agentexecution/selection_load.go @@ -0,0 +1,198 @@ +package agentexecution + +import ( + "context" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" +) + +type ExecutionHead struct { + AgentID uuid.UUID + CurrentTurnID uuid.UUID + StopSequence int64 + AnsweredThroughSequence int64 + MaxNormalInputSequence int64 + MaxContextInputSequence int64 + NormalContextID uuid.UUID + CompactionContextID uuid.UUID + PendingToolOutputID uuid.UUID + PendingOutputLimitID uuid.UUID + PendingConfigInputID uuid.UUID + PendingCheckpointID uuid.UUID + TurnContinuable bool + IncompleteTools bool + LogicalReadyAt *time.Time +} + +type ExecutionSnapshot struct { + Head ExecutionHead + View ExecutionView + Selection Selection +} + +type executionLoader struct { + q *executiondb.Queries + db executiondb.DBTX + route AgentRoute +} + +func (h *Handle) executionLoader(ctx context.Context) (*executionLoader, error) { + route, err := h.Route() + if err != nil { + return nil, err + } + l := &executionLoader{q: executiondb.New(), db: h.unit.DB(), route: route} + scope, err := l.q.LoadExecutionScope( + ctx, + l.db, + executiondb.LoadExecutionScopeParams{ProjectID: route.ProjectID, ID: route.AgentID}, + ) + if err != nil { + return nil, fmt.Errorf("load execution scope: %w", err) + } + if scope != route.RootAgentID { + return nil, fmt.Errorf("%w: agent %s root", ErrInvalidState, route.AgentID) + } + return l, nil +} + +func (h *Handle) LoadExecution(ctx context.Context) (ExecutionSnapshot, error) { + route, err := h.Route() + if err != nil { + return ExecutionSnapshot{}, err + } + var head *ExecutionHead + m := h.mutation + if m != nil { + head = &m.head + if m.loaded != nil && m.loaded.Selection.Wait != WaitModelDeadline { + return cloneExecutionSnapshot(*m.loaded), nil + } + } + snapshot, err := loadExecution(ctx, h.unit.DB(), route, head) + if err == nil && m != nil { + m.loaded = &snapshot + } + return cloneExecutionSnapshot(snapshot), err +} + +func executionHead(row executiondb.LoadExecutionBaseRow) ExecutionHead { + return ExecutionHead{ + AgentID: row.AgentID, CurrentTurnID: valueOrZero(row.CurrentTurnID), StopSequence: row.StopSequence, + AnsweredThroughSequence: row.AnsweredThroughSequence, MaxNormalInputSequence: row.MaxNormalInputSequence, + MaxContextInputSequence: row.MaxContextInputSequence, NormalContextID: valueOrZero(row.NormalContextID), + CompactionContextID: valueOrZero( + row.CompactionContextID, + ), PendingToolOutputID: valueOrZero(row.PendingToolOutputID), + PendingOutputLimitID: valueOrZero( + row.PendingOutputLimitID, + ), PendingConfigInputID: valueOrZero(row.PendingConfigInputID), + PendingCheckpointID: valueOrZero(row.PendingCheckpointID), TurnContinuable: row.TurnContinuable, + IncompleteTools: row.IncompleteTools, LogicalReadyAt: row.LogicalReadyAt, + } +} + +func valueOrZero[T any](value *T) (zero T) { + if value != nil { + return *value + } + return zero +} + +func (l *executionLoader) load(ctx context.Context, head ExecutionHead) (ExecutionSnapshot, error) { + return loadExecution(ctx, l.db, l.route, &head) +} + +func contextRecord(row executiondb.ReconstructExecutionContextsRow) ModelContext { + return ModelContext{ + ID: row.ID, + AgentID: row.AgentID, + TurnID: row.TurnID, + Operation: Operation(row.OperationKind), + Attempt: row.AttemptNumber, + InputEventSequence: row.InputEventSequence, + SourceEventSequenceEnd: valueOrZero(row.SourceEventSequenceEnd), + State: ContextState(row.State), + Recovery: RecoveryKind(valueOrZero(row.RecoveryKind)), + RetryAt: row.RetryAt, + Opening: Opening{ + InputIDs: row.OpeningInputIds, + EventSequence: row.OpeningEventSequence, + }, + } +} + +func (l *executionLoader) turn(ctx context.Context, head ExecutionHead) (*Turn, error) { + row, err := l.q.LoadExecutionTurn( + ctx, + l.db, + executiondb.LoadExecutionTurnParams{AgentID: head.AgentID, ID: head.CurrentTurnID}, + ) + if err != nil { + return nil, err + } + turn := &Turn{ + ID: row.ID, + FirstOpeningSequence: row.FirstOpeningSequence, + LastOpeningSequence: row.LastOpeningSequence, + FirstContentSequence: row.FirstContentSequence, + LatestSemantic: EventBoundary{Sequence: row.SemanticSequence, Time: row.SemanticTime}, + } + if head.MaxContextInputSequence >= turn.FirstOpeningSequence || + head.StopSequence > turn.FirstOpeningSequence { + return turn, nil + } + opening, err := l.q.CaptureExecutionOpening(ctx, l.db, executiondb.CaptureExecutionOpeningParams{ + ProjectID: l.route.ProjectID, AgentID: head.AgentID, TurnID: turn.ID, + Watermark: turn.LastOpeningSequence, StopSequence: head.StopSequence}) + if err != nil { + return nil, err + } + for i, input := range opening { + turn.InitialOpening.InputIDs = append(turn.InitialOpening.InputIDs, input.ID) + if i == 0 { + turn.InitialOpening.EventSequence = input.Sequence + turn.InitialReadyAt = input.CreatedAt + } + if input.CreatedAt.Before(turn.InitialReadyAt) { + turn.InitialReadyAt = input.CreatedAt + } + } + return turn, nil +} + +func (l *executionLoader) batch(ctx context.Context, output ModelOutput) (*ToolBatch, error) { + row, err := l.q.LoadExecutionBatch( + ctx, + l.db, + executiondb.LoadExecutionBatchParams{AgentID: l.route.AgentID, ModelOutputID: output.ID}, + ) + if err != nil || !row.HasTools { + return nil, err + } + batch := &ToolBatch{Output: output, HasIncomplete: row.Incomplete, HasRunnable: row.Runnable} + if row.Incomplete { + return batch, nil + } + completion, err := l.q.LoadExecutionBatchCompletion( + ctx, + l.db, + executiondb.LoadExecutionBatchCompletionParams{ + AgentID: l.route.AgentID, ModelOutputID: output.ID}, + ) + if err != nil { + return nil, err + } + if completion.Calls == 0 || completion.Calls != completion.Results || + completion.Calls != completion.Events { + return nil, fmt.Errorf("%w: output %s incomplete result authority", ErrInvalidState, output.ID) + } + batch.Completion = &BatchCompletion{ + LastResultSequence: completion.LastResultSequence, + ReadyAt: completion.ReadyAt, + } + return batch, nil +} diff --git a/internal/storage/internal/agentexecution/selection_test.go b/internal/storage/internal/agentexecution/selection_test.go new file mode 100644 index 000000000..6fccd6709 --- /dev/null +++ b/internal/storage/internal/agentexecution/selection_test.go @@ -0,0 +1,415 @@ +package agentexecution_test + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func identity(n byte) uuid.UUID { return uuid.UUID{15: n} } + +func decisionTime() time.Time { return time.Date(2026, 10, 9, 12, 0, 0, 0, time.UTC) } + +func baseView() agentexecution.ExecutionView { + opening := agentexecution.Opening{InputIDs: []uuid.UUID{identity(3)}, EventSequence: 10} + return agentexecution.ExecutionView{ + AgentID: identity(1), State: agentexecution.AgentActive, + MaxNormalInputSequence: 10, MaxContextInputSequence: 10, + Turn: &agentexecution.Turn{ + ID: identity(2), FirstOpeningSequence: 10, LastOpeningSequence: 10, FirstContentSequence: 10, + LatestSemantic: agentexecution.EventBoundary{Sequence: 10, Time: decisionTime().Add(-time.Minute)}, + InitialOpening: opening, InitialReadyAt: decisionTime().Add(-time.Minute), + }, + NormalContext: &agentexecution.ModelContext{ + ID: identity(4), AgentID: identity(1), TurnID: identity(2), Operation: agentexecution.OperationNormal, + Attempt: 1, InputEventSequence: 10, State: agentexecution.ContextSucceeded, Opening: opening, + }, + } +} + +func initialView() agentexecution.ExecutionView { + view := baseView() + view.NormalContext = nil + view.MaxNormalInputSequence, view.MaxContextInputSequence = 0, 0 + return view +} + +func retryView() agentexecution.ExecutionView { + view := baseView() + deadline := decisionTime().Add(-time.Second) + view.NormalContext.State, view.NormalContext.Recovery = agentexecution.ContextFailed, agentexecution.RecoveryRetry + view.NormalContext.RetryAt = &deadline + return view +} + +func boundary(view agentexecution.ExecutionView, id byte, sequence int64) *agentexecution.BoundaryWork { + return &agentexecution.BoundaryWork{ + ID: identity(id), AgentID: view.AgentID, TurnID: view.Turn.ID, + Event: agentexecution.EventBoundary{Sequence: sequence, Time: decisionTime().Add(-time.Second)}, + Opening: view.Turn.InitialOpening, + } +} + +func output(view agentexecution.ExecutionView) agentexecution.ModelOutput { + return agentexecution.ModelOutput{ + ID: identity(5), EventID: identity(6), Context: *view.NormalContext, + Event: agentexecution.EventBoundary{Sequence: 20, Time: decisionTime().Add(-time.Second)}, + } +} + +func batchView() agentexecution.ExecutionView { + view := baseView() + view.ToolBatch = &agentexecution.ToolBatch{ + Output: output(view), HasIncomplete: true, HasRunnable: true, + } + return view +} + +func completedBatchView() agentexecution.ExecutionView { + view := batchView() + view.ToolBatch.HasIncomplete, view.ToolBatch.HasRunnable = false, false + view.ToolBatch.Completion = &agentexecution.BatchCompletion{ + LastResultSequence: 30, ReadyAt: decisionTime().Add(-time.Second), + } + return view +} + +func compactionView() agentexecution.ExecutionView { + view := retryView() + view.NormalContext.Recovery, view.NormalContext.RetryAt = agentexecution.RecoveryCompact, nil + compaction := *view.NormalContext + compaction.ID, compaction.Operation = identity(7), agentexecution.OperationCompaction + compaction.SourceEventSequenceEnd = 9 + compaction.State, compaction.Recovery = agentexecution.ContextStarted, agentexecution.RecoveryNone + view.CompactionContext = &compaction + return view +} + +func TestSelectNextPrecedence(t *testing.T) { + t.Parallel() + for _, test := range []struct { + name string + view func() agentexecution.ExecutionView + change func(*agentexecution.ExecutionView) + work agentexecution.WorkKind + admit agentexecution.AdmissionKind + wait agentexecution.WaitKind + origin agentexecution.CandidateOrigin + continuable bool + incomplete bool + }{ + {name: "idle", view: baseView, wait: agentexecution.WaitIdle}, + {name: "initial", view: initialView, work: agentexecution.WorkModel, + origin: agentexecution.OriginInitial, continuable: true}, + {name: "retry", view: retryView, work: agentexecution.WorkModel, + origin: agentexecution.OriginRetry, continuable: true}, + {name: "later semantic", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.Turn.LatestSemantic = agentexecution.EventBoundary{Sequence: 15, Time: decisionTime().Add(-time.Second)} + }, work: agentexecution.WorkModel, origin: agentexecution.OriginSemantic, continuable: true}, + {name: "config", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.Config = boundary(*v, 8, 15) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginConfig, continuable: true}, + {name: "completed tools", view: completedBatchView, work: agentexecution.WorkModel, + origin: agentexecution.OriginTools, continuable: true}, + {name: "output limit", view: baseView, change: func(v *agentexecution.ExecutionView) { + o := output(*v) + v.OutputLimit = &o + }, work: agentexecution.WorkModel, origin: agentexecution.OriginOutputLimit, continuable: true}, + {name: "checkpoint", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.Checkpoint = boundary(*v, 9, 15) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginCheckpoint, continuable: true}, + {name: "all steering before queued", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.Inputs = agentexecution.InputAvailability{Steering: true, Queued: true} + }, work: agentexecution.WorkInput, admit: agentexecution.AdmitAllSteering}, + {name: "one queued", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.Inputs.Queued = true + }, work: agentexecution.WorkInput, admit: agentexecution.AdmitOneQueued}, + {name: "tools before inputs", view: batchView, change: func(v *agentexecution.ExecutionView) { + v.Inputs = agentexecution.InputAvailability{Steering: true, Queued: true} + }, work: agentexecution.WorkTool, continuable: true, incomplete: true}, + {name: "tools before config and checkpoint", view: batchView, change: func(v *agentexecution.ExecutionView) { + v.Config, v.Checkpoint = boundary(*v, 8, 21), boundary(*v, 9, 22) + }, work: agentexecution.WorkTool, continuable: true, incomplete: true}, + {name: "external tools block everything", view: batchView, change: func(v *agentexecution.ExecutionView) { + v.ToolBatch.HasRunnable = false + v.Inputs = agentexecution.InputAvailability{Steering: true, Queued: true} + v.Config, v.Checkpoint = boundary(*v, 8, 21), boundary(*v, 9, 22) + }, wait: agentexecution.WaitToolBatch, continuable: true, incomplete: true}, + {name: "steering before initial", view: initialView, change: func(v *agentexecution.ExecutionView) { + v.Inputs.Steering = true + }, work: agentexecution.WorkInput, admit: agentexecution.AdmitAllSteering, + origin: agentexecution.OriginInitial, continuable: true}, + {name: "steering retains retry candidate", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.Inputs.Steering = true + deadline := decisionTime().Add(time.Hour) + v.NormalContext.RetryAt = &deadline + }, work: agentexecution.WorkInput, admit: agentexecution.AdmitAllSteering, + origin: agentexecution.OriginRetry, continuable: true}, + {name: "model before queued", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.Inputs.Queued = true + }, work: agentexecution.WorkModel, origin: agentexecution.OriginRetry, continuable: true}, + {name: "future retry blocks queued", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.Inputs.Queued = true + deadline := decisionTime().Add(time.Hour) + v.NormalContext.RetryAt = &deadline + }, wait: agentexecution.WaitModelDeadline, origin: agentexecution.OriginRetry, continuable: true}, + {name: "initial before config", view: initialView, change: func(v *agentexecution.ExecutionView) { + v.Config = boundary(*v, 8, 15) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginInitial, continuable: true}, + {name: "semantic before config", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.Config = boundary(*v, 8, 15) + v.Turn.LatestSemantic = v.Config.Event + }, work: agentexecution.WorkModel, origin: agentexecution.OriginSemantic, continuable: true}, + {name: "config before completed batch", view: completedBatchView, change: func(v *agentexecution.ExecutionView) { + v.Config = boundary(*v, 8, 31) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginConfig, continuable: true}, + {name: "config before output limit", view: baseView, change: func(v *agentexecution.ExecutionView) { + o := output(*v) + v.OutputLimit, v.Config = &o, boundary(*v, 8, 21) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginConfig, continuable: true}, + {name: "config before checkpoint", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.Config, v.Checkpoint = boundary(*v, 8, 15), boundary(*v, 9, 16) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginConfig, continuable: true}, + {name: "completed batch before checkpoint", view: completedBatchView, change: func(v *agentexecution.ExecutionView) { + v.Checkpoint = boundary(*v, 9, 31) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginTools, continuable: true}, + {name: "future retry before completed batch and checkpoint", view: completedBatchView, + change: func(v *agentexecution.ExecutionView) { + deadline := decisionTime().Add(time.Hour) + v.NormalContext.ID, v.NormalContext.InputEventSequence = identity(12), 30 + v.MaxNormalInputSequence, v.MaxContextInputSequence = 30, 30 + v.NormalContext.State, v.NormalContext.Recovery = agentexecution.ContextFailed, agentexecution.RecoveryRetry + v.NormalContext.RetryAt = &deadline + v.Checkpoint = boundary(*v, 9, 31) + v.Inputs.Queued = true + }, wait: agentexecution.WaitModelDeadline, origin: agentexecution.OriginRetry, continuable: true}, + {name: "output limit before checkpoint", view: baseView, change: func(v *agentexecution.ExecutionView) { + o := output(*v) + v.OutputLimit, v.Checkpoint = &o, boundary(*v, 9, 21) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginOutputLimit, continuable: true}, + {name: "empty initial falls through", view: initialView, change: func(v *agentexecution.ExecutionView) { + v.Turn.InitialOpening = agentexecution.Opening{} + v.Config = boundary(*v, 8, 15) + v.Config.Opening = baseView().Turn.InitialOpening + }, work: agentexecution.WorkModel, origin: agentexecution.OriginConfig, continuable: true}, + {name: "empty config falls through", view: completedBatchView, change: func(v *agentexecution.ExecutionView) { + v.Config = boundary(*v, 8, 31) + v.Config.Opening = agentexecution.Opening{} + }, work: agentexecution.WorkModel, origin: agentexecution.OriginTools, continuable: true}, + {name: "empty retry falls through", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.NormalContext.Opening = agentexecution.Opening{} + v.Checkpoint = boundary(*v, 9, 31) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginCheckpoint, continuable: true}, + {name: "empty semantic falls through", view: retryView, change: func(v *agentexecution.ExecutionView) { + v.NormalContext.Opening = agentexecution.Opening{} + v.Config = boundary(*v, 8, 31) + v.Turn.LatestSemantic = v.Config.Event + }, work: agentexecution.WorkModel, origin: agentexecution.OriginConfig, continuable: true}, + {name: "empty tools fall through", view: completedBatchView, change: func(v *agentexecution.ExecutionView) { + v.ToolBatch.Output.Context.Opening = agentexecution.Opening{} + v.Checkpoint = boundary(*v, 9, 31) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginCheckpoint, continuable: true}, + {name: "empty output limit falls through", view: baseView, change: func(v *agentexecution.ExecutionView) { + o := output(*v) + o.Context.Opening = agentexecution.Opening{} + v.OutputLimit, v.Checkpoint = &o, boundary(*v, 9, 31) + }, work: agentexecution.WorkModel, origin: agentexecution.OriginCheckpoint, continuable: true}, + {name: "empty checkpoint falls through to queued", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.Checkpoint = boundary(*v, 9, 31) + v.Checkpoint.Opening = agentexecution.Opening{} + v.Inputs.Queued = true + }, work: agentexecution.WorkInput, admit: agentexecution.AdmitOneQueued}, + {name: "config only with unanswered prior inputs", view: initialView, change: func(v *agentexecution.ExecutionView) { + v.Turn.FirstContentSequence = 0 + v.Turn.InitialOpening.EventSequence = 2 + }, work: agentexecution.WorkModel, origin: agentexecution.OriginInitial, continuable: true}, + {name: "config only without unanswered inputs", view: initialView, change: func(v *agentexecution.ExecutionView) { + v.Turn.FirstContentSequence = 0 + v.Turn.InitialOpening = agentexecution.Opening{} + v.Config = boundary(*v, 8, 10) + }, wait: agentexecution.WaitIdle}, + {name: "started context is continuable without work", view: baseView, change: func(v *agentexecution.ExecutionView) { + v.NormalContext.State = agentexecution.ContextStarted + }, wait: agentexecution.WaitIdle, continuable: true}, + {name: "compaction suppresses normal recovery", view: compactionView, + wait: agentexecution.WaitIdle, continuable: true}, + {name: "terminal compaction suppresses normal recovery", view: compactionView, + change: func(v *agentexecution.ExecutionView) { v.CompactionContext.State = agentexecution.ContextSucceeded }, + wait: agentexecution.WaitIdle}, + {name: "compaction retry", view: compactionView, change: func(v *agentexecution.ExecutionView) { + deadline := decisionTime() + v.CompactionContext.State = agentexecution.ContextFailed + v.CompactionContext.Recovery, v.CompactionContext.RetryAt = agentexecution.RecoveryRetry, &deadline + }, work: agentexecution.WorkModel, origin: agentexecution.OriginRetry, continuable: true}, + {name: "compaction semantic restart", view: compactionView, change: func(v *agentexecution.ExecutionView) { + v.Turn.LatestSemantic.Sequence = 15 + }, work: agentexecution.WorkModel, origin: agentexecution.OriginSemantic, continuable: true}, + {name: "archived suppresses work", view: batchView, change: func(v *agentexecution.ExecutionView) { + v.State = agentexecution.AgentArchived + v.Inputs.Steering = true + }, wait: agentexecution.WaitArchived}, + {name: "no turn admits content", view: func() agentexecution.ExecutionView { + return agentexecution.ExecutionView{AgentID: identity(1), State: agentexecution.AgentActive, + Inputs: agentexecution.InputAvailability{Queued: true}} + }, work: agentexecution.WorkInput, admit: agentexecution.AdmitOneQueued}, + } { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + view := test.view() + if test.change != nil { + test.change(&view) + } + got, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, test.work, got.Work) + require.Equal(t, test.admit, got.Admission) + require.Equal(t, test.wait, got.Wait) + require.Equal(t, test.continuable, got.TurnContinuable) + require.Equal(t, test.incomplete, got.IncompleteTools) + if test.origin == 0 { + require.Nil(t, got.Model) + } else { + require.NotNil(t, got.Model) + require.Equal(t, test.origin, got.Model.Origin) + require.Equal(t, view.Turn.ID, got.Model.TurnID) + kind := agentexecution.ModelStart + if test.origin == agentexecution.OriginRetry { + kind = agentexecution.ModelResume + } else if test.origin == agentexecution.OriginTools || test.origin == agentexecution.OriginOutputLimit { + kind = agentexecution.ModelContinue + } + require.Equal(t, kind, got.Model.Kind) + } + if got.Work != agentexecution.WorkNone { + require.NotNil(t, got.LogicalReadyAt) + require.False(t, got.LogicalReadyAt.After(decisionTime())) + } else if test.wait != agentexecution.WaitModelDeadline { + require.Nil(t, got.LogicalReadyAt) + } + }) + } +} + +func TestSelectNextStrictBoundaries(t *testing.T) { + t.Parallel() + for _, kind := range []string{"initial stop", "context stop", "batch stop", "config stop", "checkpoint stop", + "output limit stop", "initial watermark", "config watermark", "checkpoint watermark", "output watermark", + "semantic watermark", "retry deadline", "config first content"} { + t.Run(kind, func(t *testing.T) { + t.Parallel() + for _, offset := range []int64{-1, 0, 1} { + view := baseView() + want := offset <= 0 + switch kind { + case "initial stop": + view = initialView() + view.StopSequence = 10 + offset + case "context stop": + view = retryView() + view.StopSequence = 10 + offset + case "batch stop": + view = completedBatchView() + view.StopSequence = 20 + offset + case "config stop": + view.Config = boundary(view, 8, 20) + view.StopSequence = 20 + offset + case "checkpoint stop": + view.Checkpoint = boundary(view, 9, 20) + view.StopSequence = 20 + offset + case "output limit stop": + o := output(view) + view.OutputLimit, view.StopSequence = &o, 20+offset + case "initial watermark": + view = initialView() + view.MaxContextInputSequence = 10 + offset + want = offset < 0 + case "config watermark": + view.Config = boundary(view, 8, 10+offset) + want = offset > 0 + case "checkpoint watermark": + view.Checkpoint = boundary(view, 9, 10+offset) + want = offset > 0 + case "output watermark": + o := output(view) + view.OutputLimit, view.NormalContext = &o, nil + view.MaxNormalInputSequence, view.MaxContextInputSequence = 20+offset, 20+offset + want = offset < 0 + case "semantic watermark": + view.NormalContext.State = agentexecution.ContextStarted + view.Turn.LatestSemantic.Sequence = 10 + offset + want = offset > 0 + case "retry deadline": + view = retryView() + deadline := decisionTime().Add(time.Duration(offset) * time.Nanosecond) + view.NormalContext.RetryAt = &deadline + case "config first content": + view.Turn.FirstContentSequence = 20 + view.Config = boundary(view, 8, 20+offset) + want = offset >= 0 + } + got, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err, "offset %d", offset) + require.Equal(t, want, got.Work == agentexecution.WorkModel, "offset %d: %+v", offset, got) + } + }) + } +} + +func TestSelectNextPreservesSourceLineage(t *testing.T) { + t.Parallel() + for _, origin := range []agentexecution.CandidateOrigin{ + agentexecution.OriginTools, agentexecution.OriginOutputLimit, agentexecution.OriginRetry, + agentexecution.OriginSemantic, agentexecution.OriginConfig, agentexecution.OriginCheckpoint, + } { + view := completedBatchView() + inherited := agentexecution.Opening{InputIDs: []uuid.UUID{identity(11), identity(3)}, EventSequence: 2} + view.NormalContext.Opening = inherited + view.ToolBatch.Output.Context.Opening = inherited + switch origin { + case agentexecution.OriginTools: + case agentexecution.OriginOutputLimit: + o := view.ToolBatch.Output + view.ToolBatch, view.OutputLimit = nil, &o + case agentexecution.OriginRetry, agentexecution.OriginSemantic: + view.ToolBatch = nil + view.NormalContext.State, view.NormalContext.Recovery = agentexecution.ContextFailed, agentexecution.RecoveryRetry + deadline := decisionTime() + view.NormalContext.RetryAt = &deadline + if origin == agentexecution.OriginSemantic { + view.Turn.LatestSemantic.Sequence = 25 + } + case agentexecution.OriginConfig: + view.ToolBatch, view.Config = nil, boundary(view, 8, 25) + view.Config.Opening = inherited + case agentexecution.OriginCheckpoint: + view.ToolBatch, view.Checkpoint = nil, boundary(view, 9, 25) + view.Checkpoint.Opening = inherited + default: + t.Fatalf("unexpected test origin %d", origin) + } + got, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.NotNil(t, got.Model) + require.Equal(t, origin, got.Model.Origin) + require.Equal(t, inherited, got.Model.Opening) + switch origin { + case agentexecution.OriginConfig: + require.Equal(t, identity(8), got.Model.SourceInputID) + case agentexecution.OriginCheckpoint: + require.Equal(t, identity(9), got.Model.SourceCheckpointID) + default: + require.Equal(t, identity(4), got.Model.SourceContextID) + } + if origin == agentexecution.OriginTools || origin == agentexecution.OriginOutputLimit { + require.Equal(t, identity(5), got.Model.SourceOutputID) + } + got.Model.Opening.InputIDs[0] = identity(99) + require.Equal(t, identity(11), inherited.InputIDs[0]) + again, err := agentexecution.SelectNext(view, decisionTime()) + require.NoError(t, err) + require.Equal(t, inherited, again.Model.Opening) + } +} diff --git a/internal/storage/internal/agentexecution/snapshot.go b/internal/storage/internal/agentexecution/snapshot.go new file mode 100644 index 000000000..328d0d6ff --- /dev/null +++ b/internal/storage/internal/agentexecution/snapshot.go @@ -0,0 +1,63 @@ +package agentexecution + +import "slices" + +func copyPointer[T any](value *T) *T { + if value == nil { + return nil + } + cloned := *value + return &cloned +} + +func cloneOpening(opening Opening) Opening { + opening.InputIDs = slices.Clone(opening.InputIDs) + return opening +} + +func cloneModel(model ModelContext) ModelContext { + model.Opening = cloneOpening(model.Opening) + model.RetryAt = copyPointer(model.RetryAt) + return model +} + +func cloneExecutionSnapshot(snapshot ExecutionSnapshot) ExecutionSnapshot { + snapshot.Head.LogicalReadyAt = copyPointer(snapshot.Head.LogicalReadyAt) + snapshot.Selection.LogicalReadyAt = copyPointer(snapshot.Selection.LogicalReadyAt) + snapshot.Selection.Model = copyPointer(snapshot.Selection.Model) + snapshot.Selection.Tool = copyPointer(snapshot.Selection.Tool) + if snapshot.Selection.Model != nil { + snapshot.Selection.Model.Opening = cloneOpening(snapshot.Selection.Model.Opening) + } + v := &snapshot.View + v.Turn = copyPointer(v.Turn) + if v.Turn != nil { + v.Turn.InitialOpening = cloneOpening(v.Turn.InitialOpening) + } + v.NormalContext = copyPointer(v.NormalContext) + if v.NormalContext != nil { + *v.NormalContext = cloneModel(*v.NormalContext) + } + v.CompactionContext = copyPointer(v.CompactionContext) + if v.CompactionContext != nil { + *v.CompactionContext = cloneModel(*v.CompactionContext) + } + v.ToolBatch = copyPointer(v.ToolBatch) + if v.ToolBatch != nil { + v.ToolBatch.Output.Context = cloneModel(v.ToolBatch.Output.Context) + v.ToolBatch.Completion = copyPointer(v.ToolBatch.Completion) + } + v.OutputLimit = copyPointer(v.OutputLimit) + if v.OutputLimit != nil { + v.OutputLimit.Context = cloneModel(v.OutputLimit.Context) + } + v.Config = copyPointer(v.Config) + if v.Config != nil { + v.Config.Opening = cloneOpening(v.Config.Opening) + } + v.Checkpoint = copyPointer(v.Checkpoint) + if v.Checkpoint != nil { + v.Checkpoint.Opening = cloneOpening(v.Checkpoint.Opening) + } + return snapshot +} diff --git a/internal/storage/internal/agentexecution/sqlc_integration_test.go b/internal/storage/internal/agentexecution/sqlc_integration_test.go new file mode 100644 index 000000000..1efd04e0d --- /dev/null +++ b/internal/storage/internal/agentexecution/sqlc_integration_test.go @@ -0,0 +1,26 @@ +//go:build integration + +package agentexecution_test + +import ( + "net/url" + "os" + "os/exec" + "testing" + + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + "github.com/stretchr/testify/require" +) + +func TestExecutionSQLCPrepare(t *testing.T) { + pool := integrationdb.OpenMigratedPool(t, t.Context(), "../../../../migrations") + command := exec.CommandContext(t.Context(), "go", "tool", "-modfile=tools/ci/go.mod", + "github.com/sqlc-dev/sqlc/cmd/sqlc", "vet", "-f", "sqlc.vet-db.yaml") + command.Dir = "../../../.." + databaseURL, err := url.Parse(os.Getenv("OMNARA_TEST_DATABASE_URL")) + require.NoError(t, err) + databaseURL.Path = "/" + pool.Config().ConnConfig.Database + command.Env = append(os.Environ(), "SQLC_DATABASE_URL="+databaseURL.String()) + output, err := command.CombinedOutput() + require.NoError(t, err, string(output)) +} diff --git a/internal/storage/internal/agentexecution/state.go b/internal/storage/internal/agentexecution/state.go new file mode 100644 index 000000000..6a865b2f2 --- /dev/null +++ b/internal/storage/internal/agentexecution/state.go @@ -0,0 +1,198 @@ +package agentexecution + +import ( + "time" + + "github.com/google/uuid" +) + +type AgentState string + +const ( + AgentActive AgentState = "active" + AgentArchived AgentState = "archived" +) + +type Operation string + +const ( + OperationNormal Operation = "normal" + OperationCompaction Operation = "compaction" +) + +type ContextState string + +const ( + ContextStarted ContextState = "started" + ContextSucceeded ContextState = "succeeded" + ContextFailed ContextState = "failed" + ContextCanceled ContextState = "canceled" +) + +type RecoveryKind string + +const ( + RecoveryNone RecoveryKind = "" + RecoveryRetry RecoveryKind = "retry" + RecoveryCompact RecoveryKind = "compact" + RecoveryReduce RecoveryKind = "reduce_compaction_source" +) + +type Opening struct { + InputIDs []uuid.UUID + EventSequence int64 +} + +type EventBoundary struct { + Sequence int64 + Time time.Time +} + +type Turn struct { + ID uuid.UUID + FirstOpeningSequence int64 + LastOpeningSequence int64 + FirstContentSequence int64 + LatestSemantic EventBoundary + InitialOpening Opening + InitialReadyAt time.Time +} + +type ModelContext struct { + ID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + Operation Operation + Attempt int32 + InputEventSequence int64 + SourceEventSequenceEnd int64 + State ContextState + Recovery RecoveryKind + RetryAt *time.Time + Opening Opening +} + +type ModelOutput struct { + ID uuid.UUID + EventID uuid.UUID + Event EventBoundary + Context ModelContext +} + +type BatchCompletion struct { + LastResultSequence int64 + ReadyAt time.Time +} + +type ToolBatch struct { + Output ModelOutput + HasIncomplete bool + HasRunnable bool + Completion *BatchCompletion +} + +type BoundaryWork struct { + ID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + Event EventBoundary + Opening Opening +} + +type InputAvailability struct { + Steering bool + Queued bool +} + +type ExecutionView struct { + AgentID uuid.UUID + State AgentState + Turn *Turn + StopSequence int64 + MaxNormalInputSequence int64 + MaxContextInputSequence int64 + NormalContext *ModelContext + CompactionContext *ModelContext + ToolBatch *ToolBatch + OutputLimit *ModelOutput + Config *BoundaryWork + Checkpoint *BoundaryWork + Inputs InputAvailability +} + +type WorkKind uint8 + +const ( + WorkNone WorkKind = iota + WorkTool + WorkModel + WorkInput +) + +type AdmissionKind uint8 + +const ( + AdmitNone AdmissionKind = iota + AdmitAllSteering + AdmitOneQueued +) + +type WaitKind uint8 + +const ( + WaitNone WaitKind = iota + WaitIdle + WaitToolBatch + WaitModelDeadline + WaitArchived +) + +type ModelKind string + +const ( + ModelStart ModelKind = "start" + ModelResume ModelKind = "resume" + ModelContinue ModelKind = "continue" +) + +type CandidateOrigin uint8 + +const ( + OriginInitial CandidateOrigin = iota + 1 + OriginRetry + OriginSemantic + OriginConfig + OriginTools + OriginOutputLimit + OriginCheckpoint +) + +type ModelDecision struct { + Kind ModelKind + Origin CandidateOrigin + TurnID uuid.UUID + SourceContextID uuid.UUID + SourceOutputID uuid.UUID + SourceInputID uuid.UUID + SourceCheckpointID uuid.UUID + Opening Opening + ReadyAt time.Time +} + +type ToolDecision struct { + TurnID uuid.UUID + SourceContextID uuid.UUID + OutputID uuid.UUID + SourceEventID uuid.UUID +} + +type Selection struct { + Work WorkKind + Admission AdmissionKind + Wait WaitKind + Model *ModelDecision + Tool *ToolDecision + TurnContinuable bool + IncompleteTools bool + LogicalReadyAt *time.Time +} diff --git a/internal/storage/internal/agentexecution/unit.go b/internal/storage/internal/agentexecution/unit.go new file mode 100644 index 000000000..48f8c524c --- /dev/null +++ b/internal/storage/internal/agentexecution/unit.go @@ -0,0 +1,189 @@ +package agentexecution + +import ( + "context" + "errors" + "fmt" + "maps" + "slices" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/notifications" +) + +var ErrUnitClosed = errors.New("agent execution unit is closed") +var ErrUnitAborted = errors.New("agent execution unit requires rollback") +var ErrLockPlan = errors.New("agent execution lock plan must be restarted") + +type Unit struct { + cell *AgentCell + tx pgx.Tx + notifications *notifications.TxNotifications + handles map[uuid.UUID]*Handle + webhooks map[uuid.UUID]webhookTarget + exclusiveProjects map[uuid.UUID]struct{} + runtimeLocked bool + closed bool + aborted bool + finalizing bool + savepoints int +} + +func newUnit(cell *AgentCell, tx pgx.Tx) *Unit { + return &Unit{ + cell: cell, tx: tx, notifications: notifications.NewTxNotifications(), + handles: make(map[uuid.UUID]*Handle), webhooks: make(map[uuid.UUID]webhookTarget), + exclusiveProjects: make(map[uuid.UUID]struct{}), + } +} + +func (u *Unit) DB() unitDB { return unitDB{unit: u} } + +func (u *Unit) Notifications() *notifications.TxNotifications { return u.notifications } + +func (u *Unit) CellID() CellID { return u.cell.id } + +func (u *Unit) active() error { + if u.closed { + return ErrUnitClosed + } + if u.aborted { + return ErrUnitAborted + } + return nil +} + +func (u *Unit) Commit(ctx context.Context, operation string) (err error) { + if err := u.active(); err != nil { + return err + } + if u.savepoints != 0 || u.finalizing { + return errors.New("cannot commit inside unit savepoint or finalization") + } + u.finalizing = true + defer func() { + if err != nil { + u.aborted = true + err = fmt.Errorf("commit %s: %w", operation, err) + } + }() + if err := u.publishExecution(ctx); err != nil { + return err + } + if err := u.enqueueWebhooks(ctx); err != nil { + return err + } + if err := u.tx.Commit(ctx); err != nil { + u.closed = true + return err + } + u.closed = true + u.notifications.Flush(context.WithoutCancel(ctx), u.cell.publisher) + return nil +} + +func (u *Unit) Rollback(ctx context.Context) error { + if u.closed { + return pgx.ErrTxClosed + } + if u.savepoints != 0 { + return errors.New("use savepoint rollback inside a unit savepoint") + } + u.closed = true + return u.tx.Rollback(ctx) +} + +func (u *Unit) orderedHandles() []*Handle { + handles := slices.Collect(maps.Values(u.handles)) + slices.SortFunc(handles, func(a, b *Handle) int { return compareRoutes(a.route, b.route) }) + return handles +} + +func (u *Unit) Savepoint(ctx context.Context, run func(*Unit) error) (err error) { + if err := u.active(); err != nil { + return err + } + if u.finalizing { + return errors.New("cannot create savepoint during finalization") + } + parent := u.tx + nested, err := parent.Begin(ctx) + if err != nil { + return err + } + snapshot := u.snapshot() + u.tx = nested + u.savepoints++ + released := false + defer func() { + u.savepoints-- + u.tx = parent + if !released { + rollbackErr := nested.Rollback(context.WithoutCancel(ctx)) + u.restore(snapshot) + if rollbackErr != nil { + u.aborted = true + err = errors.Join(err, fmt.Errorf("rollback agent execution savepoint: %w", rollbackErr)) + } + } + }() + if err := run(u); err != nil { + return err + } + if err := u.active(); err != nil { + return err + } + if err := nested.Commit(ctx); err != nil { + return err + } + released = true + return nil +} + +type unitSnapshot struct { + handles map[uuid.UUID]*Handle + states map[uuid.UUID]Handle + notifications *notifications.TxNotifications + webhooks map[uuid.UUID]webhookTarget + exclusiveProjects map[uuid.UUID]struct{} + runtimeLocked bool +} + +func (u *Unit) snapshot() unitSnapshot { + s := unitSnapshot{ + handles: maps.Clone(u.handles), states: make(map[uuid.UUID]Handle), + notifications: u.notifications.Clone(), webhooks: maps.Clone(u.webhooks), + exclusiveProjects: maps.Clone(u.exclusiveProjects), + runtimeLocked: u.runtimeLocked, + } + for id, h := range u.handles { + state := *h + if h.mutation != nil { + state.mutation = h.mutation.clone() + } + s.states[id] = state + } + for id, target := range s.webhooks { + target.events = slices.Clone(target.events) + s.webhooks[id] = target + } + return s +} + +func (u *Unit) restore(snapshot unitSnapshot) { + for id, h := range u.handles { + if snapshot.handles[id] != h { + h.valid = false + } + } + u.handles = snapshot.handles + for id, h := range u.handles { + *h = snapshot.states[id] + } + u.notifications.Restore(snapshot.notifications) + u.webhooks = snapshot.webhooks + u.exclusiveProjects = snapshot.exclusiveProjects + u.runtimeLocked = snapshot.runtimeLocked + u.aborted = false +} diff --git a/internal/storage/internal/agentexecution/unit_integration_test.go b/internal/storage/internal/agentexecution/unit_integration_test.go new file mode 100644 index 000000000..ecb3db301 --- /dev/null +++ b/internal/storage/internal/agentexecution/unit_integration_test.go @@ -0,0 +1,272 @@ +//go:build integration + +package agentexecution + +import ( + "context" + "errors" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" + "github.com/stretchr/testify/require" +) + +type unitPublisher struct { + publish func(context.Context, notifications.PostCommitIntent) +} + +func (p unitPublisher) PublishPostCommit(ctx context.Context, intent notifications.PostCommitIntent) { + p.publish(ctx, intent) +} + +func unitTestCell(t *testing.T, publisher notifications.PostCommitPublisher) (*AgentCell, *pgxpool.Pool) { + t.Helper() + pool := integrationdb.OpenUnmigratedPool(t, t.Context()) + _, err := pool.Exec(t.Context(), ` +CREATE TABLE agents(id uuid PRIMARY KEY, project_id uuid NOT NULL, org_id uuid NOT NULL, + parent_agent_id uuid, root_agent_id uuid NOT NULL, state text NOT NULL DEFAULT 'active'); +CREATE TABLE unit_probe(id bigint GENERATED ALWAYS AS IDENTITY, agent_id uuid, phase text, labels text[]); +CREATE TABLE deferred_probe(id integer UNIQUE DEFERRABLE INITIALLY DEFERRED);`) + require.NoError(t, err) + return NewCell("test", pool, publisher, nil), pool +} + +func unitTestAgent(t *testing.T, u *Unit, project, agent uuid.UUID) *Handle { + t.Helper() + require.NoError(t, u.CreatedAgent(project, agent, uuid.Nil)) + h, err := u.Agent(AgentRoute{CellID: u.CellID(), ProjectID: project, AgentID: agent, RootAgentID: agent}) + require.NoError(t, err) + return h +} + +func TestUnitCommitLifecycle(t *testing.T) { + ctx := t.Context() + var published []notifications.PostCommitIntent + var pool *pgxpool.Pool + cell, pool := unitTestCell( + t, + unitPublisher{publish: func(ctx context.Context, intent notifications.PostCommitIntent) { + require.NoError(t, ctx.Err()) + require.Nil(t, ctx.Done()) + var phases []string + require.NoError( + t, + pool.QueryRow(ctx, `SELECT array_agg(phase ORDER BY id) FROM unit_probe`).Scan(&phases), + ) + require.Equal(t, []string{"write"}, phases) + published = append(published, intent) + }}, + ) + u, err := cell.Begin(ctx) + require.NoError(t, err) + defer func() { _ = u.Rollback(ctx) }() + _, raw := any(u.DB()).(pgx.Tx) + require.False(t, raw) + rows, err := u.DB().Query(ctx, `SELECT 1`) + require.NoError(t, err) + require.Nil(t, rows.Conn()) + rows.Close() + require.NoError(t, rows.Err()) + project := uuid.New() + higher := uuid.MustParse("ffffffff-ffff-ffff-ffff-ffffffffffff") + lower := uuid.MustParse("00000000-0000-0000-0000-000000000001") + unitTestAgent(t, u, project, higher) + unitTestAgent(t, u, project, lower) + require.Equal(t, lower, u.orderedHandles()[0].route.AgentID) + _, err = u.DB().Exec(ctx, `INSERT INTO unit_probe(agent_id,phase) VALUES($1,'write')`, lower) + require.NoError(t, err) + u.Notifications().AddDaemonWork(uuid.New()) + require.Empty(t, published) + require.NoError(t, u.Commit(ctx, "probe")) + require.Len(t, published, 1) + var agents []uuid.UUID + require.NoError( + t, + pool.QueryRow(ctx, `SELECT array_agg(agent_id ORDER BY id) FROM unit_probe`).Scan(&agents), + ) + require.Equal(t, []uuid.UUID{lower}, agents) + _, err = u.handles[lower].Route() + require.ErrorIs(t, err, ErrUnitClosed) + require.ErrorIs(t, u.Commit(ctx, "again"), ErrUnitClosed) +} + +func TestUnitSavepointSnapshots(t *testing.T) { + ctx := t.Context() + var published []notifications.PostCommitIntent + cell, pool := unitTestCell( + t, + unitPublisher{publish: func(_ context.Context, intent notifications.PostCommitIntent) { + published = append(published, intent) + }}, + ) + u, err := cell.Begin(ctx) + require.NoError(t, err) + defer func() { _ = u.Rollback(ctx) }() + project, agent, machine := uuid.New(), uuid.New(), uuid.New() + h := unitTestAgent(t, u, project, agent) + mutation := &executionMutation{head: ExecutionHead{StopSequence: 1}} + h.mutation = mutation + pending := u.Notifications() + pending.AddDaemonWork(machine) + process := uuid.New() + pending.AddDaemonProcessTermination(machine, process) + u.webhooks[agent] = webhookTarget{orgID: uuid.New(), events: []string{"model_output"}} + originalTarget := u.webhooks[agent] + rejected := errors.New("reject inner") + var escaped *Handle + require.ErrorIs(t, u.Savepoint(ctx, func(u *Unit) error { + mutation.head.StopSequence = 2 + pending.AddDaemonWork(uuid.New()) + pending.AddDaemonProcessTermination(machine, uuid.New()) + u.webhooks[agent] = webhookTarget{orgID: uuid.New()} + delete(u.handles, agent) + escaped = unitTestAgent(t, u, project, uuid.New()) + escaped.mutation = &executionMutation{head: ExecutionHead{StopSequence: 3}} + delete(u.handles, escaped.route.AgentID) + _, err := u.DB(). + Exec(ctx, `INSERT INTO unit_probe(phase) VALUES ('rollback'); SELECT pg_advisory_xact_lock(8173)`) + require.NoError(t, err) + require.NoError(t, u.Savepoint(ctx, func(u *Unit) error { + _, err := u.DB().Exec(ctx, `INSERT INTO unit_probe(phase) VALUES ('nested release')`) + return err + })) + require.Empty(t, published) + return rejected + }), rejected) + require.Same(t, pending, u.Notifications()) + require.Same(t, h, u.handles[agent]) + restored := h.mutation + require.EqualValues(t, 1, restored.head.StopSequence) + require.Equal(t, originalTarget, u.webhooks[agent]) + _, err = escaped.Route() + require.ErrorIs(t, err, ErrUnitClosed) + var unlocked bool + require.NoError(t, pool.QueryRow(ctx, `SELECT pg_try_advisory_xact_lock(8173)`).Scan(&unlocked)) + require.True(t, unlocked) + require.NoError(t, u.Savepoint(ctx, func(u *Unit) error { + require.Error(t, u.Commit(ctx, "nested")) + require.Error(t, u.Rollback(ctx)) + restored.head.StopSequence = 4 + return nil + })) + require.NoError(t, u.Commit(ctx, "snapshot")) + require.Len(t, published, 2) + for _, intent := range published { + if terminate, ok := intent.(notifications.DaemonProcessTerminationCommitted); ok { + require.Equal(t, []uuid.UUID{process}, terminate.ProcessIDs) + } + } + var phases []string + require.NoError( + t, + pool.QueryRow(ctx, `SELECT array_agg(phase ORDER BY id) FROM unit_probe`).Scan(&phases), + ) + require.Empty(t, phases) +} + +func TestUnitFailuresSuppressPublication(t *testing.T) { + for _, scenario := range []string{"rollback", "deferred commit", "aborted SQL", "cancel"} { + t.Run(scenario, func(t *testing.T) { + ctx := t.Context() + cell, pool := unitTestCell( + t, + unitPublisher{publish: func(context.Context, notifications.PostCommitIntent) { + t.Error("notification from failed transaction") + }}, + ) + u, err := cell.Begin(ctx) + require.NoError(t, err) + defer func() { _ = u.Rollback(ctx) }() + unitTestAgent(t, u, uuid.New(), uuid.New()) + u.Notifications().AddDaemonWork(uuid.New()) + switch scenario { + case "rollback": + require.NoError(t, u.Rollback(ctx)) + case "deferred commit": + _, err := u.DB().Exec(ctx, `INSERT INTO deferred_probe VALUES (1),(1)`) + require.NoError(t, err) + case "aborted SQL": + _, err := u.DB().Exec(ctx, `SELECT 1 / 0`) + require.Error(t, err) + _, err = u.DB().Exec(ctx, `SELECT 1`) + require.ErrorIs(t, err, ErrUnitAborted) + case "cancel": + canceled, cancel := context.WithCancel(ctx) + cancel() + ctx = canceled + } + require.Error(t, u.Commit(ctx, scenario)) + _ = u.Rollback(t.Context()) + var count int + require.NoError(t, pool.QueryRow(t.Context(), `SELECT count(*) FROM unit_probe`).Scan(&count)) + require.Zero(t, count) + }) + } +} + +func TestUnitSavepointRecoversDatabaseErrors(t *testing.T) { + ctx := t.Context() + cell, _ := unitTestCell(t, nil) + require.NoError(t, cell.Transact(ctx, func(u *Unit) error { + err := u.Savepoint(ctx, func(u *Unit) error { + _, err := u.DB().Exec(ctx, `SELECT 1 / 0`) + require.Error(t, err) + return nil + }) + require.ErrorIs(t, err, ErrUnitAborted) + _, err = u.DB().Exec(ctx, `INSERT INTO deferred_probe VALUES (1)`) + require.NoError(t, err) + return err + })) +} + +func TestUnitLockPlansAndSavepointLocks(t *testing.T) { + ctx := t.Context() + cell, pool := unitTestCell(t, nil) + project, org := uuid.New(), uuid.New() + lower := uuid.MustParse("00000000-0000-0000-0000-000000000001") + higher := uuid.MustParse("ffffffff-ffff-ffff-ffff-ffffffffffff") + _, err := pool.Exec( + ctx, + `INSERT INTO agents(id,project_id,org_id,parent_agent_id,root_agent_id) VALUES ($1,$3,$4,NULL,$1),($2,$3,$4,$1,$1)`, + lower, + higher, + project, + org, + ) + require.NoError(t, err) + u, err := cell.Begin(ctx) + require.NoError(t, err) + defer func() { _ = u.Rollback(ctx) }() + refs := []lifecyclelock.AgentRef{{ProjectID: project, AgentID: higher}} + require.NoError(t, u.LockAgentRefs(ctx, refs, LifecycleAuthority{})) + require.Equal(t, lower, u.handles[higher].route.RootAgentID) + require.ErrorIs(t, u.LockAgentRefs( + ctx, + []lifecyclelock.AgentRef{{ProjectID: project, AgentID: lower}}, IngressAuthority{}), + ErrLockPlan, + ) + plan, err := u.PlanAgents(ctx, refs, LifecycleAuthority{}) + require.NoError(t, err) + other, err := NewCell("other", pool, nil, nil).Begin(ctx) + require.NoError(t, err) + defer func() { _ = other.Rollback(ctx) }() + require.ErrorIs(t, other.LockAgents(ctx, plan), ErrLockPlan) + rollback := errors.New("rollback locks") + require.ErrorIs(t, other.Savepoint(ctx, func(u *Unit) error { + return errors.Join(u.LockAgentRefs( + ctx, + []lifecyclelock.AgentRef{{ProjectID: project, AgentID: lower}}, RepairAuthority{}), + rollback, + ) + }), rollback) + require.Empty(t, other.handles) + check := integrationdb.BeginTx(t, ctx, pool) + _, err = check.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, lower) + require.NoError(t, err) +} diff --git a/internal/storage/internal/agentexecution/webhooks.go b/internal/storage/internal/agentexecution/webhooks.go new file mode 100644 index 000000000..4f6006420 --- /dev/null +++ b/internal/storage/internal/agentexecution/webhooks.go @@ -0,0 +1,112 @@ +package agentexecution + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "slices" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" +) + +type webhookTarget struct { + orgID uuid.UUID + events []string +} + +func (u *Unit) CaptureWebhookEligibility(ctx context.Context, agents ...uuid.UUID) error { + ids := make(map[uuid.UUID]struct{}) + for _, id := range agents { + ids[id] = struct{}{} + } + for id := range u.notifications.AgentEvents() { + ids[id] = struct{}{} + } + for _, update := range u.notifications.ToolCallUpdates() { + ids[update.AgentID] = struct{}{} + } + q := dbsqlc.New(u.DB()) + for id := range ids { + if _, captured := u.webhooks[id]; captured { + continue + } + row, err := q.GetAgentEventWebhookTarget(ctx, dbsqlc.GetAgentEventWebhookTargetParams{ID: id}) + if errors.Is(err, pgx.ErrNoRows) { + u.webhooks[id] = webhookTarget{} + continue + } + if err != nil { + return fmt.Errorf("load event webhook target: %w", err) + } + if row.Url == "" { + u.webhooks[id] = webhookTarget{} + continue + } + var events []string + if err := json.Unmarshal(row.Events, &events); err != nil { + return fmt.Errorf("decode event webhook events: %w", err) + } + u.webhooks[id] = webhookTarget{orgID: row.OrgID, events: events} + } + return nil +} + +func webhookJSON(value json.RawMessage) *json.RawMessage { + if len(value) == 0 { + return nil + } + return &value +} + +func (u *Unit) enqueueWebhooks(ctx context.Context) error { + pending := u.notifications + if pending == nil || (len(pending.AgentEvents()) == 0 && len(pending.ToolCallUpdates()) == 0) { + return nil + } + if err := u.CaptureWebhookEligibility(ctx); err != nil { + return err + } + q := dbsqlc.New(u.DB()) + webhookTargets := u.webhooks + for agentID, events := range pending.AgentEvents() { + target, enabled := webhookTargets[agentID] + if !enabled || target.orgID == uuid.Nil { + continue + } + for _, event := range events { + if !slices.Contains(target.events, event.Kind) { + continue + } + if err := q.EnqueueEventWebhookDelivery(ctx, dbsqlc.EnqueueEventWebhookDeliveryParams{ + OrgID: target.orgID, AgentID: agentID, EventSequence: &event.Sequence, + }); err != nil { + return fmt.Errorf("enqueue event webhook: %w", err) + } + } + } + for _, update := range pending.ToolCallUpdates() { + target, enabled := webhookTargets[update.AgentID] + if !enabled || !slices.Contains(target.events, "tool_call_update") { + continue + } + var interactionUpdate json.RawMessage + if update.InteractionUpdate != nil { + var err error + interactionUpdate, err = json.Marshal(update.InteractionUpdate) + if err != nil { + return fmt.Errorf("encode webhook interaction update: %w", err) + } + } + if err := q.EnqueueEventWebhookDelivery(ctx, dbsqlc.EnqueueEventWebhookDeliveryParams{ + OrgID: target.orgID, AgentID: update.AgentID, + ToolCallID: &update.ToolCallID, ToolState: &update.State, + InteractionUpdate: webhookJSON(interactionUpdate), + }); err != nil { + return fmt.Errorf("enqueue tool webhook: %w", err) + } + } + return nil +} diff --git a/internal/storage/internal/dbsqlc/agent_inputs.sql.go b/internal/storage/internal/dbsqlc/agent_inputs.sql.go index 8f512af76..fd14a3a9b 100644 --- a/internal/storage/internal/dbsqlc/agent_inputs.sql.go +++ b/internal/storage/internal/dbsqlc/agent_inputs.sql.go @@ -13,121 +13,6 @@ import ( "github.com/google/uuid" ) -const admitAgentInput = `-- name: AdmitAgentInput :one -UPDATE agent_inputs input -SET state = 'resolved', - admitted_event_id = $1::uuid, - admitted_at = event.created_at, - resolved_at = event.created_at -FROM agent_events event -WHERE input.project_id = $2 - AND input.agent_id = $3 - AND input.id = $4 - AND input.state = 'received' - AND event.agent_id = input.agent_id - AND event.agent_input_id = input.id - AND event.id = $1 -RETURNING input.admitted_at, input.resolved_at -` - -type AdmitAgentInputParams struct { - AdmittedEventID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -type AdmitAgentInputRow struct { - AdmittedAt *time.Time - ResolvedAt *time.Time -} - -func (q *Queries) AdmitAgentInput(ctx context.Context, arg AdmitAgentInputParams) (AdmitAgentInputRow, error) { - row := q.db.QueryRow(ctx, admitAgentInput, - arg.AdmittedEventID, - arg.ProjectID, - arg.AgentID, - arg.ID, - ) - var i AdmitAgentInputRow - err := row.Scan(&i.AdmittedAt, &i.ResolvedAt) - return i, err -} - -const cancelQueuedBacklogInput = `-- name: CancelQueuedBacklogInput :execrows -UPDATE agent_inputs -SET state = 'canceled', - canceled_at = statement_timestamp() -WHERE agent_inputs.project_id = $1 - AND agent_inputs.agent_id = $2 - AND agent_inputs.id = $3 - AND agent_inputs.state = 'received' - AND agent_inputs.delivery_mode = 'queued' - AND agent_inputs.input_kind = 'content' -` - -type CancelQueuedBacklogInputParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) CancelQueuedBacklogInput(ctx context.Context, arg CancelQueuedBacklogInputParams) (int64, error) { - result, err := q.db.Exec(ctx, cancelQueuedBacklogInput, arg.ProjectID, arg.AgentID, arg.ID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const cancelQueuedBacklogInputsForAgent = `-- name: CancelQueuedBacklogInputsForAgent :execrows -UPDATE agent_inputs -SET state = 'canceled', - canceled_at = statement_timestamp() -WHERE project_id = $1 - AND agent_id = $2 - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' -` - -type CancelQueuedBacklogInputsForAgentParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) CancelQueuedBacklogInputsForAgent(ctx context.Context, arg CancelQueuedBacklogInputsForAgentParams) (int64, error) { - result, err := q.db.Exec(ctx, cancelQueuedBacklogInputsForAgent, arg.ProjectID, arg.AgentID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const cancelSteeringAgentInputsForAgent = `-- name: CancelSteeringAgentInputsForAgent :execrows -UPDATE agent_inputs -SET state = 'canceled', - canceled_at = statement_timestamp() -WHERE project_id = $1 - AND agent_id = $2 - AND state = 'received' - AND delivery_mode = 'steering' - AND input_kind = 'content' -` - -type CancelSteeringAgentInputsForAgentParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) CancelSteeringAgentInputsForAgent(ctx context.Context, arg CancelSteeringAgentInputsForAgentParams) (int64, error) { - result, err := q.db.Exec(ctx, cancelSteeringAgentInputsForAgent, arg.ProjectID, arg.AgentID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - const claimNextAgentWakeup = `-- name: ClaimNextAgentWakeup :one WITH locked_agent AS MATERIALIZED ( SELECT agent.id AS agent_id, agent.project_id @@ -175,165 +60,6 @@ func (q *Queries) ClaimNextAgentWakeup(ctx context.Context) (ClaimNextAgentWakeu return i, err } -const consumeAgentWakeup = `-- name: ConsumeAgentWakeup :execrows -DELETE FROM agent_wakeups wake -USING agents agent -WHERE agent.project_id = $1 - AND agent.id = $2 - AND wake.agent_id = agent.id -` - -type ConsumeAgentWakeupParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) ConsumeAgentWakeup(ctx context.Context, arg ConsumeAgentWakeupParams) (int64, error) { - result, err := q.db.Exec(ctx, consumeAgentWakeup, arg.ProjectID, arg.AgentID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const deleteAgentWakeup = `-- name: DeleteAgentWakeup :execrows -WITH locked AS ( - SELECT agent.id - FROM agents agent - WHERE agent.project_id = $1 - AND agent.id = $2 - FOR UPDATE -) -DELETE FROM agent_wakeups wake -USING locked -WHERE wake.agent_id = locked.id -` - -type DeleteAgentWakeupParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) DeleteAgentWakeup(ctx context.Context, arg DeleteAgentWakeupParams) (int64, error) { - result, err := q.db.Exec(ctx, deleteAgentWakeup, arg.ProjectID, arg.AgentID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const demoteSteeringInputToQueued = `-- name: DemoteSteeringInputToQueued :execrows -UPDATE agent_inputs input -SET delivery_mode = 'queued', - input_rank = coalesce( - ( - SELECT max(existing.input_rank) + $1::bigint - FROM agent_inputs existing - WHERE existing.project_id = input.project_id - AND existing.agent_id = input.agent_id - AND existing.delivery_mode = 'queued' - AND existing.state = 'received' - AND existing.input_kind = 'content' - ), - $1::bigint - ) -WHERE input.project_id = $2 - AND input.agent_id = $3 - AND input.id = $4 - AND input.state = 'received' - AND input.delivery_mode = 'steering' - AND input.input_kind = 'content' -` - -type DemoteSteeringInputToQueuedParams struct { - RankStride int64 - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) DemoteSteeringInputToQueued(ctx context.Context, arg DemoteSteeringInputToQueuedParams) (int64, error) { - result, err := q.db.Exec(ctx, demoteSteeringInputToQueued, - arg.RankStride, - arg.ProjectID, - arg.AgentID, - arg.ID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const getNextQueuedAgentInputForAdmission = `-- name: GetNextQueuedAgentInputForAdmission :one -SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata -FROM agent_inputs input -WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' -ORDER BY input.input_rank ASC, input.queued_at ASC, input.id ASC -FOR UPDATE -LIMIT 1 -` - -type GetNextQueuedAgentInputForAdmissionParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type GetNextQueuedAgentInputForAdmissionRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - State string - InputRank int64 - ActorID *uuid.UUID - InputKind string - IntegrationTargetID *uuid.UUID - IdempotencyScope string - InputIdempotencyKey string - QueuedAt time.Time - AdmittedEventID *uuid.UUID - AdmittedAt *time.Time - CanceledAt *time.Time - DeliveryMode string - ControlType string - TargetInteractionID *uuid.UUID - ResolvedAt *time.Time - RejectedReason string - Metadata json.RawMessage -} - -func (q *Queries) GetNextQueuedAgentInputForAdmission(ctx context.Context, arg GetNextQueuedAgentInputForAdmissionParams) (GetNextQueuedAgentInputForAdmissionRow, error) { - row := q.db.QueryRow(ctx, getNextQueuedAgentInputForAdmission, arg.ProjectID, arg.AgentID) - var i GetNextQueuedAgentInputForAdmissionRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.State, - &i.InputRank, - &i.ActorID, - &i.InputKind, - &i.IntegrationTargetID, - &i.IdempotencyScope, - &i.InputIdempotencyKey, - &i.QueuedAt, - &i.AdmittedEventID, - &i.AdmittedAt, - &i.CanceledAt, - &i.DeliveryMode, - &i.ControlType, - &i.TargetInteractionID, - &i.ResolvedAt, - &i.RejectedReason, - &i.Metadata, - ) - return i, err -} - const listQueuedBacklogInputs = `-- name: ListQueuedBacklogInputs :many SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata FROM agent_inputs input @@ -435,542 +161,3 @@ func (q *Queries) ListQueuedBacklogInputs(ctx context.Context, arg ListQueuedBac } return items, nil } - -const listSteeringAgentInputsForAdmission = `-- name: ListSteeringAgentInputsForAdmission :many -SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata -FROM agent_inputs input -WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.state = 'received' - AND input.delivery_mode = 'steering' - AND input.input_kind = 'content' - AND NOT agent_has_incomplete_tool_batch($1, $2) -ORDER BY input.input_rank ASC, input.queued_at ASC, input.id ASC -FOR UPDATE -` - -type ListSteeringAgentInputsForAdmissionParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type ListSteeringAgentInputsForAdmissionRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - State string - InputRank int64 - ActorID *uuid.UUID - InputKind string - IntegrationTargetID *uuid.UUID - IdempotencyScope string - InputIdempotencyKey string - QueuedAt time.Time - AdmittedEventID *uuid.UUID - AdmittedAt *time.Time - CanceledAt *time.Time - DeliveryMode string - ControlType string - TargetInteractionID *uuid.UUID - ResolvedAt *time.Time - RejectedReason string - Metadata json.RawMessage -} - -func (q *Queries) ListSteeringAgentInputsForAdmission(ctx context.Context, arg ListSteeringAgentInputsForAdmissionParams) ([]ListSteeringAgentInputsForAdmissionRow, error) { - rows, err := q.db.Query(ctx, listSteeringAgentInputsForAdmission, arg.ProjectID, arg.AgentID) - if err != nil { - return nil, err - } - defer rows.Close() - items := []ListSteeringAgentInputsForAdmissionRow{} - for rows.Next() { - var i ListSteeringAgentInputsForAdmissionRow - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.State, - &i.InputRank, - &i.ActorID, - &i.InputKind, - &i.IntegrationTargetID, - &i.IdempotencyScope, - &i.InputIdempotencyKey, - &i.QueuedAt, - &i.AdmittedEventID, - &i.AdmittedAt, - &i.CanceledAt, - &i.DeliveryMode, - &i.ControlType, - &i.TargetInteractionID, - &i.ResolvedAt, - &i.RejectedReason, - &i.Metadata, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const markAgentWakeup = `-- name: MarkAgentWakeup :exec -INSERT INTO agent_wakeups(agent_id, ready_at, updated_at, metadata) -SELECT agent.id, - statement_timestamp(), statement_timestamp(), $1 -FROM agents agent -WHERE agent.project_id = $2 - AND agent.id = $3 - AND agent.state <> 'archived' -ON CONFLICT (agent_id) DO UPDATE -SET ready_at = LEAST(agent_wakeups.ready_at, excluded.ready_at), - updated_at = GREATEST(agent_wakeups.updated_at, excluded.updated_at), - metadata = agent_wakeups.metadata || excluded.metadata -` - -type MarkAgentWakeupParams struct { - Metadata json.RawMessage - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) MarkAgentWakeup(ctx context.Context, arg MarkAgentWakeupParams) error { - _, err := q.db.Exec(ctx, markAgentWakeup, arg.Metadata, arg.ProjectID, arg.AgentID) - return err -} - -const moveQueuedBacklogInputAfter = `-- name: MoveQueuedBacklogInputAfter :execrows -WITH target AS ( - SELECT anchor.id, anchor.input_rank, anchor.queued_at - FROM agent_inputs anchor - WHERE anchor.project_id = $1 - AND anchor.agent_id = $2 - AND anchor.state = 'received' - AND anchor.delivery_mode = 'queued' - AND anchor.input_kind = 'content' - AND anchor.id = $4 -), -next_neighbor AS ( - SELECT candidate.input_rank - FROM agent_inputs candidate - CROSS JOIN target - WHERE candidate.project_id = $1 - AND candidate.agent_id = $2 - AND candidate.state = 'received' - AND candidate.delivery_mode = 'queued' - AND candidate.input_kind = 'content' - AND candidate.id <> $3 - AND (candidate.input_rank, candidate.queued_at, candidate.id) > (target.input_rank, target.queued_at, target.id) - ORDER BY candidate.input_rank ASC, candidate.queued_at ASC, candidate.id ASC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT id FROM target) IS NULL THEN NULL - WHEN (SELECT input_rank FROM next_neighbor) IS NULL THEN - (SELECT input_rank + $5::bigint FROM target) - WHEN (SELECT input_rank FROM next_neighbor) - (SELECT input_rank FROM target) > 1 THEN - (SELECT input_rank FROM target) + - ((SELECT input_rank FROM next_neighbor) - (SELECT input_rank FROM target)) / 2 - ELSE NULL::bigint - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL -` - -type MoveQueuedBacklogInputAfterParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - AnchorID uuid.UUID - RankStride int64 -} - -func (q *Queries) MoveQueuedBacklogInputAfter(ctx context.Context, arg MoveQueuedBacklogInputAfterParams) (int64, error) { - result, err := q.db.Exec(ctx, moveQueuedBacklogInputAfter, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.AnchorID, - arg.RankStride, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const moveQueuedBacklogInputBefore = `-- name: MoveQueuedBacklogInputBefore :execrows -WITH target AS ( - SELECT anchor.id, anchor.input_rank, anchor.queued_at - FROM agent_inputs anchor - WHERE anchor.project_id = $1 - AND anchor.agent_id = $2 - AND anchor.state = 'received' - AND anchor.delivery_mode = 'queued' - AND anchor.input_kind = 'content' - AND anchor.id = $4 -), -previous_neighbor AS ( - SELECT candidate.input_rank - FROM agent_inputs candidate - CROSS JOIN target - WHERE candidate.project_id = $1 - AND candidate.agent_id = $2 - AND candidate.state = 'received' - AND candidate.delivery_mode = 'queued' - AND candidate.input_kind = 'content' - AND candidate.id <> $3 - AND (candidate.input_rank, candidate.queued_at, candidate.id) < (target.input_rank, target.queued_at, target.id) - ORDER BY candidate.input_rank DESC, candidate.queued_at DESC, candidate.id DESC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT id FROM target) IS NULL THEN NULL - WHEN (SELECT input_rank FROM previous_neighbor) IS NULL AND - (SELECT input_rank FROM target) > 1 THEN - (SELECT input_rank / 2 FROM target) - WHEN (SELECT input_rank FROM target) - (SELECT input_rank FROM previous_neighbor) > 1 THEN - (SELECT input_rank FROM previous_neighbor) + - ((SELECT input_rank FROM target) - (SELECT input_rank FROM previous_neighbor)) / 2 - ELSE NULL::bigint - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL -` - -type MoveQueuedBacklogInputBeforeParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - AnchorID uuid.UUID -} - -func (q *Queries) MoveQueuedBacklogInputBefore(ctx context.Context, arg MoveQueuedBacklogInputBeforeParams) (int64, error) { - result, err := q.db.Exec(ctx, moveQueuedBacklogInputBefore, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.AnchorID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const moveQueuedBacklogInputToBack = `-- name: MoveQueuedBacklogInputToBack :execrows -WITH last_neighbor AS ( - SELECT input_rank - FROM agent_inputs - WHERE project_id = $1 - AND agent_id = $2 - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' - AND id <> $3 - ORDER BY input_rank DESC, queued_at DESC, id DESC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT input_rank FROM last_neighbor) IS NULL THEN $4::bigint - ELSE (SELECT input_rank + $4::bigint FROM last_neighbor) - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL -` - -type MoveQueuedBacklogInputToBackParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - RankStride int64 -} - -func (q *Queries) MoveQueuedBacklogInputToBack(ctx context.Context, arg MoveQueuedBacklogInputToBackParams) (int64, error) { - result, err := q.db.Exec(ctx, moveQueuedBacklogInputToBack, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.RankStride, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const moveQueuedBacklogInputToFront = `-- name: MoveQueuedBacklogInputToFront :execrows -WITH first_neighbor AS ( - SELECT input_rank - FROM agent_inputs - WHERE project_id = $1 - AND agent_id = $2 - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' - AND id <> $3 - ORDER BY input_rank ASC, queued_at ASC, id ASC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT input_rank FROM first_neighbor) IS NULL THEN $4::bigint - WHEN (SELECT input_rank FROM first_neighbor) > 1 THEN - (SELECT input_rank / 2 FROM first_neighbor) - ELSE NULL::bigint - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL -` - -type MoveQueuedBacklogInputToFrontParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - RankStride int64 -} - -func (q *Queries) MoveQueuedBacklogInputToFront(ctx context.Context, arg MoveQueuedBacklogInputToFrontParams) (int64, error) { - result, err := q.db.Exec(ctx, moveQueuedBacklogInputToFront, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.RankStride, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const promoteQueuedInputToSteering = `-- name: PromoteQueuedInputToSteering :one -WITH promoted AS ( - UPDATE agent_inputs input - SET delivery_mode = 'steering', - input_rank = coalesce( - ( - SELECT max(existing.input_rank) + $4::bigint - FROM agent_inputs existing - WHERE existing.project_id = input.project_id - AND existing.agent_id = input.agent_id - AND existing.delivery_mode = 'steering' - AND existing.state = 'received' - AND existing.input_kind = 'content' - ), - $4::bigint - ) - WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - RETURNING TRUE -) -SELECT EXISTS (SELECT 1 FROM promoted) AS changed, - EXISTS (SELECT 1 FROM promoted) - OR EXISTS ( - SELECT 1 - FROM agent_inputs input - WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.input_kind = 'content' - AND ( - (input.state = 'received' AND input.delivery_mode = 'steering') - OR (input.state = 'resolved' AND input.admitted_event_id IS NOT NULL) - ) - ) AS effective -` - -type PromoteQueuedInputToSteeringParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - RankStride int64 -} - -type PromoteQueuedInputToSteeringRow struct { - Changed bool - Effective *bool -} - -func (q *Queries) PromoteQueuedInputToSteering(ctx context.Context, arg PromoteQueuedInputToSteeringParams) (PromoteQueuedInputToSteeringRow, error) { - row := q.db.QueryRow(ctx, promoteQueuedInputToSteering, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.RankStride, - ) - var i PromoteQueuedInputToSteeringRow - err := row.Scan(&i.Changed, &i.Effective) - return i, err -} - -const queuedBacklogMoveIsValid = `-- name: QueuedBacklogMoveIsValid :one -SELECT CASE WHEN EXISTS ( - SELECT 1 - FROM agent_inputs input - WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' -) AND ( - NOT $4::boolean - OR EXISTS ( - SELECT 1 - FROM agent_inputs anchor - WHERE anchor.project_id = $1 - AND anchor.agent_id = $2 - AND anchor.id = $5 - AND anchor.state = 'received' - AND anchor.delivery_mode = 'queued' - AND anchor.input_kind = 'content' - ) -) THEN true ELSE false END AS valid -` - -type QueuedBacklogMoveIsValidParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - RequiresAnchor bool - AnchorID uuid.UUID -} - -func (q *Queries) QueuedBacklogMoveIsValid(ctx context.Context, arg QueuedBacklogMoveIsValidParams) (bool, error) { - row := q.db.QueryRow(ctx, queuedBacklogMoveIsValid, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.RequiresAnchor, - arg.AnchorID, - ) - var valid bool - err := row.Scan(&valid) - return valid, err -} - -const rebalanceQueuedBacklogRanks = `-- name: RebalanceQueuedBacklogRanks :execrows -WITH ordered AS ( - SELECT id, row_number() OVER (ORDER BY input_rank ASC, queued_at ASC, id ASC) AS row_number - FROM agent_inputs - WHERE project_id = $2 - AND agent_id = $3 - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' -) -UPDATE agent_inputs input -SET input_rank = ordered.row_number * $1::bigint -FROM ordered -WHERE input.project_id = $2 - AND input.agent_id = $3 - AND input.id = ordered.id -` - -type RebalanceQueuedBacklogRanksParams struct { - RankStride int64 - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) RebalanceQueuedBacklogRanks(ctx context.Context, arg RebalanceQueuedBacklogRanksParams) (int64, error) { - result, err := q.db.Exec(ctx, rebalanceQueuedBacklogRanks, arg.RankStride, arg.ProjectID, arg.AgentID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const reconcileAgentWakeup = `-- name: ReconcileAgentWakeup :exec -WITH locked_agent AS MATERIALIZED ( - SELECT agent.id, agent.project_id - FROM agents agent - WHERE agent.project_id = $1 - AND agent.id = $2 - AND agent.state <> 'archived' -), -desired AS MATERIALIZED ( - SELECT agent_next_wakeup_ready_at( - locked_agent.project_id, - locked_agent.id - ) AS ready_at - FROM locked_agent -), -upserted AS ( - INSERT INTO agent_wakeups(agent_id, ready_at, updated_at, metadata) - SELECT locked_agent.id, - desired.ready_at, - statement_timestamp(), - $3 - FROM locked_agent - JOIN desired ON desired.ready_at IS NOT NULL - ON CONFLICT (agent_id) DO UPDATE - SET ready_at = excluded.ready_at, - updated_at = GREATEST(agent_wakeups.updated_at, excluded.updated_at), - metadata = agent_wakeups.metadata || excluded.metadata - RETURNING agent_id -) -DELETE FROM agent_wakeups wake -USING locked_agent -CROSS JOIN desired -WHERE wake.agent_id = locked_agent.id - AND desired.ready_at IS NULL -` - -type ReconcileAgentWakeupParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - Metadata json.RawMessage -} - -func (q *Queries) ReconcileAgentWakeup(ctx context.Context, arg ReconcileAgentWakeupParams) error { - _, err := q.db.Exec(ctx, reconcileAgentWakeup, arg.ProjectID, arg.AgentID, arg.Metadata) - return err -} diff --git a/internal/storage/internal/dbsqlc/agent_inputs_ingress.sql.go b/internal/storage/internal/dbsqlc/agent_inputs_ingress.sql.go index ce1cb6fa5..e786f527f 100644 --- a/internal/storage/internal/dbsqlc/agent_inputs_ingress.sql.go +++ b/internal/storage/internal/dbsqlc/agent_inputs_ingress.sql.go @@ -13,23 +13,21 @@ import ( "github.com/google/uuid" ) -const getAgentInputByIdempotency = `-- name: GetAgentInputByIdempotency :one +const getAgentInput = `-- name: GetAgentInput :one SELECT id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata FROM agent_inputs WHERE project_id = $1 AND agent_id = $2 - AND idempotency_scope = $3::text - AND input_idempotency_key = $4::text + AND id = $3 ` -type GetAgentInputByIdempotencyParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - IdempotencyScope string - InputIdempotencyKey string +type GetAgentInputParams struct { + ProjectID uuid.UUID + AgentID uuid.UUID + ID uuid.UUID } -type GetAgentInputByIdempotencyRow struct { +type GetAgentInputRow struct { ID uuid.UUID ProjectID uuid.UUID AgentID uuid.UUID @@ -53,14 +51,9 @@ type GetAgentInputByIdempotencyRow struct { Metadata json.RawMessage } -func (q *Queries) GetAgentInputByIdempotency(ctx context.Context, arg GetAgentInputByIdempotencyParams) (GetAgentInputByIdempotencyRow, error) { - row := q.db.QueryRow(ctx, getAgentInputByIdempotency, - arg.ProjectID, - arg.AgentID, - arg.IdempotencyScope, - arg.InputIdempotencyKey, - ) - var i GetAgentInputByIdempotencyRow +func (q *Queries) GetAgentInput(ctx context.Context, arg GetAgentInputParams) (GetAgentInputRow, error) { + row := q.db.QueryRow(ctx, getAgentInput, arg.ProjectID, arg.AgentID, arg.ID) + var i GetAgentInputRow err := row.Scan( &i.ID, &i.ProjectID, @@ -87,49 +80,23 @@ func (q *Queries) GetAgentInputByIdempotency(ctx context.Context, arg GetAgentIn return i, err } -const insertAgentInput = `-- name: InsertAgentInput :one -WITH generated AS ( - SELECT coalesce($10, uuidv7()) AS id -) -INSERT INTO agent_inputs(id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, delivery_mode, idempotency_scope, input_idempotency_key, queued_at, metadata) -SELECT generated.id, agent.project_id, agent.id, 'received', - coalesce( - ( - SELECT max(existing.input_rank) + $1::bigint - FROM agent_inputs existing - WHERE existing.project_id = agent.project_id - AND existing.agent_id = agent.id - AND existing.delivery_mode = $2::text - AND existing.state = 'received' - AND existing.input_kind = 'content' - ), - $1::bigint - ), - $3, 'content', - $4::uuid, - $2::text, - $5, $6, statement_timestamp(), $7 -FROM agents agent -JOIN generated ON true -WHERE agent.project_id = $8 - AND agent.id = $9 -RETURNING id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata +const getAgentInputByIdempotency = `-- name: GetAgentInputByIdempotency :one +SELECT id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata +FROM agent_inputs +WHERE project_id = $1 + AND agent_id = $2 + AND idempotency_scope = $3::text + AND input_idempotency_key = $4::text ` -type InsertAgentInputParams struct { - RankStride int64 - DeliveryMode string - ActorID *uuid.UUID - IntegrationTargetID *uuid.UUID - IdempotencyScope *string - InputIdempotencyKey *string - Metadata json.RawMessage +type GetAgentInputByIdempotencyParams struct { ProjectID uuid.UUID AgentID uuid.UUID - ID *uuid.UUID + IdempotencyScope string + InputIdempotencyKey string } -type InsertAgentInputRow struct { +type GetAgentInputByIdempotencyRow struct { ID uuid.UUID ProjectID uuid.UUID AgentID uuid.UUID @@ -153,20 +120,14 @@ type InsertAgentInputRow struct { Metadata json.RawMessage } -func (q *Queries) InsertAgentInput(ctx context.Context, arg InsertAgentInputParams) (InsertAgentInputRow, error) { - row := q.db.QueryRow(ctx, insertAgentInput, - arg.RankStride, - arg.DeliveryMode, - arg.ActorID, - arg.IntegrationTargetID, - arg.IdempotencyScope, - arg.InputIdempotencyKey, - arg.Metadata, +func (q *Queries) GetAgentInputByIdempotency(ctx context.Context, arg GetAgentInputByIdempotencyParams) (GetAgentInputByIdempotencyRow, error) { + row := q.db.QueryRow(ctx, getAgentInputByIdempotency, arg.ProjectID, arg.AgentID, - arg.ID, + arg.IdempotencyScope, + arg.InputIdempotencyKey, ) - var i InsertAgentInputRow + var i GetAgentInputByIdempotencyRow err := row.Scan( &i.ID, &i.ProjectID, diff --git a/internal/storage/internal/dbsqlc/agent_interactions.sql.go b/internal/storage/internal/dbsqlc/agent_interactions.sql.go index 3646fe6d6..0ae76bae2 100644 --- a/internal/storage/internal/dbsqlc/agent_interactions.sql.go +++ b/internal/storage/internal/dbsqlc/agent_interactions.sql.go @@ -7,146 +7,10 @@ package dbsqlc import ( "context" - "encoding/json" - "time" "github.com/google/uuid" ) -const agentInteractionHasLaterStop = `-- name: AgentInteractionHasLaterStop :one -SELECT EXISTS ( - SELECT 1 -FROM agent_interactions interaction - JOIN tool_call_read_projection tool_call ON tool_call.agent_id = interaction.agent_id - AND tool_call.id = interaction.tool_call_id - JOIN agent_stop_events stop_event ON stop_event.project_id = tool_call.project_id - AND stop_event.agent_id = tool_call.agent_id - AND stop_event.sequence > tool_call.source_event_sequence - WHERE tool_call.project_id = $1 - AND interaction.agent_id = $2 - AND interaction.id = $3 -) AS stopped -` - -type AgentInteractionHasLaterStopParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) AgentInteractionHasLaterStop(ctx context.Context, arg AgentInteractionHasLaterStopParams) (bool, error) { - row := q.db.QueryRow(ctx, agentInteractionHasLaterStop, arg.ProjectID, arg.AgentID, arg.ID) - var stopped bool - err := row.Scan(&stopped) - return stopped, err -} - -const cancelOpenAgentInteractionsForAgent = `-- name: CancelOpenAgentInteractionsForAgent :many -UPDATE agent_interactions interaction -SET state = 'canceled', - resolution = jsonb_build_object('reason', $1::text), - resolved_by_input_id = $2::uuid, - resolved_at = statement_timestamp() -FROM agent_interaction_read_projection projection -WHERE interaction.agent_id = $3 - AND interaction.state = 'open' - AND projection.project_id = $4 - AND projection.agent_id = interaction.agent_id - AND projection.id = interaction.id - AND projection.turn_id = $5 -RETURNING interaction.id -` - -type CancelOpenAgentInteractionsForAgentParams struct { - Reason string - ResolvedByInputID *uuid.UUID - AgentID uuid.UUID - ProjectID uuid.UUID - TurnID uuid.UUID -} - -func (q *Queries) CancelOpenAgentInteractionsForAgent(ctx context.Context, arg CancelOpenAgentInteractionsForAgentParams) ([]uuid.UUID, error) { - rows, err := q.db.Query(ctx, cancelOpenAgentInteractionsForAgent, - arg.Reason, - arg.ResolvedByInputID, - arg.AgentID, - arg.ProjectID, - arg.TurnID, - ) - if err != nil { - return nil, err - } - defer rows.Close() - items := []uuid.UUID{} - for rows.Next() { - var id uuid.UUID - if err := rows.Scan(&id); err != nil { - return nil, err - } - items = append(items, id) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const cancelOpenAgentInteractionsForToolCall = `-- name: CancelOpenAgentInteractionsForToolCall :many -UPDATE agent_interactions -SET state = 'canceled', - resolution = jsonb_build_object('reason', $1::text), - resolved_by_input_id = NULL, - resolved_at = statement_timestamp() -WHERE agent_interactions.agent_id = $2 - AND agent_interactions.tool_call_id = $3 - AND agent_interactions.state = 'open' - AND EXISTS ( - SELECT 1 - FROM agents agent - WHERE agent.id = agent_interactions.agent_id - AND agent.project_id = $4 - ) -RETURNING id, interaction_kind, state -` - -type CancelOpenAgentInteractionsForToolCallParams struct { - Reason string - AgentID uuid.UUID - ToolCallID uuid.UUID - ProjectID uuid.UUID -} - -type CancelOpenAgentInteractionsForToolCallRow struct { - ID uuid.UUID - InteractionKind string - State string -} - -func (q *Queries) CancelOpenAgentInteractionsForToolCall(ctx context.Context, arg CancelOpenAgentInteractionsForToolCallParams) ([]CancelOpenAgentInteractionsForToolCallRow, error) { - rows, err := q.db.Query(ctx, cancelOpenAgentInteractionsForToolCall, - arg.Reason, - arg.AgentID, - arg.ToolCallID, - arg.ProjectID, - ) - if err != nil { - return nil, err - } - defer rows.Close() - items := []CancelOpenAgentInteractionsForToolCallRow{} - for rows.Next() { - var i CancelOpenAgentInteractionsForToolCallRow - if err := rows.Scan(&i.ID, &i.InteractionKind, &i.State); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const getAgentInteraction = `-- name: GetAgentInteraction :one SELECT id, project_id, agent_id, turn_id, model_call_context_id, tool_call_id, provider_call_id, @@ -235,310 +99,3 @@ func (q *Queries) GetAgentInteractionByToolCallKind(ctx context.Context, arg Get ) return i, err } - -const insertAgentInteraction = `-- name: InsertAgentInteraction :one -INSERT INTO agent_interactions( - agent_id, tool_call_id, - interaction_kind, state, request, created_at, destination -) -SELECT tool_call.agent_id, - tool_call.id, $1, 'open', - $2, statement_timestamp(), $3::jsonb -FROM tool_call_read_projection tool_call -WHERE tool_call.project_id = $4 - AND tool_call.agent_id = $5 - AND tool_call.id = $6 -RETURNING id -` - -type InsertAgentInteractionParams struct { - InteractionKind string - Request json.RawMessage - Destination *json.RawMessage - ProjectID uuid.UUID - AgentID uuid.UUID - ToolCallID uuid.UUID -} - -func (q *Queries) InsertAgentInteraction(ctx context.Context, arg InsertAgentInteractionParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, insertAgentInteraction, - arg.InteractionKind, - arg.Request, - arg.Destination, - arg.ProjectID, - arg.AgentID, - arg.ToolCallID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const insertInteractionResponseAgentInput = `-- name: InsertInteractionResponseAgentInput :one -WITH target AS ( - SELECT projection.project_id, interaction.agent_id, interaction.id - FROM agent_interactions interaction - JOIN agent_interaction_read_projection projection - ON projection.agent_id = interaction.agent_id - AND projection.id = interaction.id - WHERE projection.project_id = $5 - AND interaction.agent_id = $6 - AND interaction.id = $7 -) -INSERT INTO agent_inputs( - project_id, agent_id, state, - actor_id, input_kind, delivery_mode, target_interaction_id, - idempotency_scope, input_idempotency_key, queued_at, metadata -) -SELECT target.project_id, target.agent_id, 'received', - $1::uuid, - 'interaction_response', 'immediate', target.id, - $2, $3, - statement_timestamp(), $4 -FROM target -ON CONFLICT (project_id, agent_id, idempotency_scope, input_idempotency_key) WHERE idempotency_scope IS NOT NULL AND input_idempotency_key IS NOT NULL DO NOTHING -RETURNING id, project_id, agent_id, state, input_rank, actor_id, input_kind, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata -` - -type InsertInteractionResponseAgentInputParams struct { - ActorID *uuid.UUID - IdempotencyScope *string - InputIdempotencyKey *string - Metadata json.RawMessage - ProjectID uuid.UUID - AgentID uuid.UUID - TargetInteractionID uuid.UUID -} - -type InsertInteractionResponseAgentInputRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - State string - InputRank int64 - ActorID *uuid.UUID - InputKind string - IdempotencyScope string - InputIdempotencyKey string - QueuedAt time.Time - AdmittedEventID *uuid.UUID - AdmittedAt *time.Time - CanceledAt *time.Time - DeliveryMode string - ControlType string - TargetInteractionID *uuid.UUID - AgentConfigID *uuid.UUID - ResolvedAt *time.Time - RejectedReason string - Metadata json.RawMessage -} - -func (q *Queries) InsertInteractionResponseAgentInput(ctx context.Context, arg InsertInteractionResponseAgentInputParams) (InsertInteractionResponseAgentInputRow, error) { - row := q.db.QueryRow(ctx, insertInteractionResponseAgentInput, - arg.ActorID, - arg.IdempotencyScope, - arg.InputIdempotencyKey, - arg.Metadata, - arg.ProjectID, - arg.AgentID, - arg.TargetInteractionID, - ) - var i InsertInteractionResponseAgentInputRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.State, - &i.InputRank, - &i.ActorID, - &i.InputKind, - &i.IdempotencyScope, - &i.InputIdempotencyKey, - &i.QueuedAt, - &i.AdmittedEventID, - &i.AdmittedAt, - &i.CanceledAt, - &i.DeliveryMode, - &i.ControlType, - &i.TargetInteractionID, - &i.AgentConfigID, - &i.ResolvedAt, - &i.RejectedReason, - &i.Metadata, - ) - return i, err -} - -const listAgentInteractionsByIDs = `-- name: ListAgentInteractionsByIDs :many -SELECT id, project_id, agent_id, turn_id, - model_call_context_id, tool_call_id, provider_call_id, - interaction_kind, state, request, resolution, - resolved_by_input_id, created_at, resolved_at, destination, presentation_receipt -FROM agent_interaction_read_projection -WHERE project_id = $1 - AND agent_id = $2 - AND id = ANY($3::uuid[]) -ORDER BY created_at, id -` - -type ListAgentInteractionsByIDsParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - Ids []uuid.UUID -} - -func (q *Queries) ListAgentInteractionsByIDs(ctx context.Context, arg ListAgentInteractionsByIDsParams) ([]AgentInteractionReadProjection, error) { - rows, err := q.db.Query(ctx, listAgentInteractionsByIDs, arg.ProjectID, arg.AgentID, arg.Ids) - if err != nil { - return nil, err - } - defer rows.Close() - items := []AgentInteractionReadProjection{} - for rows.Next() { - var i AgentInteractionReadProjection - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.ModelCallContextID, - &i.ToolCallID, - &i.ProviderCallID, - &i.InteractionKind, - &i.State, - &i.Request, - &i.Resolution, - &i.ResolvedByInputID, - &i.CreatedAt, - &i.ResolvedAt, - &i.Destination, - &i.PresentationReceipt, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const resolveAgentInteraction = `-- name: ResolveAgentInteraction :one -UPDATE agent_interactions interaction -SET state = 'resolved', - resolution = $1, - resolved_by_input_id = $2::uuid, - resolved_at = statement_timestamp() -WHERE interaction.agent_id = $3 - AND interaction.id = $4 - AND interaction.state = 'open' - AND EXISTS ( - SELECT 1 - FROM agents agent - WHERE agent.id = interaction.agent_id - AND agent.project_id = $5 - ) -RETURNING interaction.id -` - -type ResolveAgentInteractionParams struct { - Resolution json.RawMessage - ResolvedByInputID *uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - ProjectID uuid.UUID -} - -func (q *Queries) ResolveAgentInteraction(ctx context.Context, arg ResolveAgentInteractionParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, resolveAgentInteraction, - arg.Resolution, - arg.ResolvedByInputID, - arg.AgentID, - arg.ID, - arg.ProjectID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const resolveInteractionResponseAgentInput = `-- name: ResolveInteractionResponseAgentInput :one -UPDATE agent_inputs -SET state = 'resolved', - admitted_event_id = $1, - admitted_at = statement_timestamp(), - resolved_at = statement_timestamp() -WHERE project_id = $2 - AND agent_id = $3 - AND id = $4 - AND input_kind = 'interaction_response' - AND target_interaction_id = $5 - AND state = 'received' -RETURNING id, project_id, agent_id, state, input_rank, actor_id, input_kind, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata -` - -type ResolveInteractionResponseAgentInputParams struct { - EventID *uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - TargetInteractionID *uuid.UUID -} - -type ResolveInteractionResponseAgentInputRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - State string - InputRank int64 - ActorID *uuid.UUID - InputKind string - IdempotencyScope string - InputIdempotencyKey string - QueuedAt time.Time - AdmittedEventID *uuid.UUID - AdmittedAt *time.Time - CanceledAt *time.Time - DeliveryMode string - ControlType string - TargetInteractionID *uuid.UUID - AgentConfigID *uuid.UUID - ResolvedAt *time.Time - RejectedReason string - Metadata json.RawMessage -} - -func (q *Queries) ResolveInteractionResponseAgentInput(ctx context.Context, arg ResolveInteractionResponseAgentInputParams) (ResolveInteractionResponseAgentInputRow, error) { - row := q.db.QueryRow(ctx, resolveInteractionResponseAgentInput, - arg.EventID, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.TargetInteractionID, - ) - var i ResolveInteractionResponseAgentInputRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.State, - &i.InputRank, - &i.ActorID, - &i.InputKind, - &i.IdempotencyScope, - &i.InputIdempotencyKey, - &i.QueuedAt, - &i.AdmittedEventID, - &i.AdmittedAt, - &i.CanceledAt, - &i.DeliveryMode, - &i.ControlType, - &i.TargetInteractionID, - &i.AgentConfigID, - &i.ResolvedAt, - &i.RejectedReason, - &i.Metadata, - ) - return i, err -} diff --git a/internal/storage/internal/dbsqlc/agent_machine_bindings.sql.go b/internal/storage/internal/dbsqlc/agent_machine_bindings.sql.go index 64fb3d12f..8d6ce126b 100644 --- a/internal/storage/internal/dbsqlc/agent_machine_bindings.sql.go +++ b/internal/storage/internal/dbsqlc/agent_machine_bindings.sql.go @@ -608,32 +608,6 @@ func (q *Queries) LockAgentMachineSources(ctx context.Context, arg LockAgentMach return err } -const lockAttachedAgentPoolMachines = `-- name: LockAttachedAgentPoolMachines :exec -SELECT machine.id -FROM machines machine -JOIN agent_machine_bindings binding ON binding.org_id = machine.org_id - AND binding.machine_id = machine.id -WHERE binding.project_id = $1 - AND binding.agent_id = $2 - AND binding.binding_kind = 'pool' - AND binding.state = 'attached' - AND machine.source_kind = 'pool' - AND machine.deleted_at IS NULL - AND machine.lifecycle_state NOT IN ('deleting', 'delete_failed', 'deleted') -ORDER BY machine.id -FOR UPDATE OF machine -` - -type LockAttachedAgentPoolMachinesParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) LockAttachedAgentPoolMachines(ctx context.Context, arg LockAttachedAgentPoolMachinesParams) error { - _, err := q.db.Exec(ctx, lockAttachedAgentPoolMachines, arg.ProjectID, arg.AgentID) - return err -} - const markAgentMachineBindingDeleteRequested = `-- name: MarkAgentMachineBindingDeleteRequested :one UPDATE agent_machine_bindings SET delete_tool_call_id = $1::uuid, diff --git a/internal/storage/internal/dbsqlc/agent_runtime.sql.go b/internal/storage/internal/dbsqlc/agent_runtime.sql.go index f4bd8a0f9..097d48e8a 100644 --- a/internal/storage/internal/dbsqlc/agent_runtime.sql.go +++ b/internal/storage/internal/dbsqlc/agent_runtime.sql.go @@ -32,33 +32,6 @@ func (q *Queries) AgentExistsInProject(ctx context.Context, arg AgentExistsInPro return exists, err } -const archiveAgent = `-- name: ArchiveAgent :execrows -UPDATE agents -SET state = 'archived', - archived_at = coalesce(archived_at, statement_timestamp()), - updated_at = CASE - WHEN state = 'archived' THEN updated_at - ELSE statement_timestamp() - END -WHERE project_id = $1 - AND id = $2 -` - -type ArchiveAgentParams struct { - ProjectID uuid.UUID - ID uuid.UUID -} - -// Idempotent: archiving an already-archived agent matches the row and keeps -// the original archived_at. -func (q *Queries) ArchiveAgent(ctx context.Context, arg ArchiveAgentParams) (int64, error) { - result, err := q.db.Exec(ctx, archiveAgent, arg.ProjectID, arg.ID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - const getAgent = `-- name: GetAgent :one SELECT id, org_id, project_id, state, name, agent_profile_id, current_config_id, interaction_target_id, @@ -261,121 +234,6 @@ func (q *Queries) GetAgentInProject(ctx context.Context, arg GetAgentInProjectPa return i, err } -const insertAgent = `-- name: InsertAgent :one -WITH inserted AS ( - INSERT INTO agents( - id, org_id, project_id, state, name, agent_profile_id, current_config_id, - idempotency_key, parent_agent_id, subagent_key, - archive_after_idle_minutes, created_at, updated_at - ) - SELECT - coalesce($1::uuid, uuidv7()), $2, $3, 'active', $4, - $5, $6, $7, - $8, $9, - $10, - transaction_timestamp(), transaction_timestamp() - FROM projects project - JOIN orgs org ON org.id = project.org_id - WHERE project.org_id = $2 - AND project.id = $3 - AND project.deleted_at IS NULL - AND org.deleted_at IS NULL - ON CONFLICT (project_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING - RETURNING id, org_id, project_id, state, name, - agent_profile_id, current_config_id, interaction_target_id, - idempotency_key, next_event_sequence, created_at, updated_at, archived_at, - parent_agent_id, subagent_key -) -SELECT agent.id, agent.org_id, agent.project_id, agent.state, agent.name, - agent.agent_profile_id, agent.current_config_id, agent.interaction_target_id, - coalesce(agent.idempotency_key, '') AS idempotency_key, - agent.next_event_sequence, agent.created_at, agent.updated_at, agent.archived_at, - agent.parent_agent_id, agent.subagent_key, - coalesce(configured_model.name, '') AS model_name, - coalesce(model_provider_config.name, '') AS model_provider_config_name -FROM inserted agent -LEFT JOIN agent_configs agent_config - ON agent_config.project_id = agent.project_id - AND agent_config.id = agent.current_config_id -LEFT JOIN configured_models configured_model - ON configured_model.org_id = agent.org_id - AND configured_model.id = agent_config.configured_model_id -LEFT JOIN model_provider_configs model_provider_config - ON model_provider_config.org_id = configured_model.org_id - AND model_provider_config.id = configured_model.model_provider_config_id -` - -type InsertAgentParams struct { - ID *uuid.UUID - OrgID uuid.UUID - ProjectID uuid.UUID - Name string - AgentProfileID *uuid.UUID - CurrentConfigID uuid.UUID - IdempotencyKey *string - ParentAgentID *uuid.UUID - SubagentKey string - ArchiveAfterIdleMinutes *int32 -} - -type InsertAgentRow struct { - ID uuid.UUID - OrgID uuid.UUID - ProjectID uuid.UUID - State string - Name string - AgentProfileID *uuid.UUID - CurrentConfigID uuid.UUID - InteractionTargetID *uuid.UUID - IdempotencyKey string - NextEventSequence int64 - CreatedAt time.Time - UpdatedAt time.Time - ArchivedAt *time.Time - ParentAgentID *uuid.UUID - SubagentKey string - ModelName string - ModelProviderConfigName string -} - -// @sqlc-vet-disable configured-models-deleted-at model-provider-configs-deleted-at -// Display-only model names must still resolve after the model or provider config is soft deleted. -func (q *Queries) InsertAgent(ctx context.Context, arg InsertAgentParams) (InsertAgentRow, error) { - row := q.db.QueryRow(ctx, insertAgent, - arg.ID, - arg.OrgID, - arg.ProjectID, - arg.Name, - arg.AgentProfileID, - arg.CurrentConfigID, - arg.IdempotencyKey, - arg.ParentAgentID, - arg.SubagentKey, - arg.ArchiveAfterIdleMinutes, - ) - var i InsertAgentRow - err := row.Scan( - &i.ID, - &i.OrgID, - &i.ProjectID, - &i.State, - &i.Name, - &i.AgentProfileID, - &i.CurrentConfigID, - &i.InteractionTargetID, - &i.IdempotencyKey, - &i.NextEventSequence, - &i.CreatedAt, - &i.UpdatedAt, - &i.ArchivedAt, - &i.ParentAgentID, - &i.SubagentKey, - &i.ModelName, - &i.ModelProviderConfigName, - ) - return i, err -} - const listAgentsForProjects = `-- name: ListAgentsForProjects :many WITH listed AS ( SELECT agent.id, diff --git a/internal/storage/internal/dbsqlc/agent_runtime_locks.sql.go b/internal/storage/internal/dbsqlc/agent_runtime_locks.sql.go index 4e8cf1b14..f8ef9b9ee 100644 --- a/internal/storage/internal/dbsqlc/agent_runtime_locks.sql.go +++ b/internal/storage/internal/dbsqlc/agent_runtime_locks.sql.go @@ -7,86 +7,11 @@ package dbsqlc import ( "context" - "encoding/json" "time" "github.com/google/uuid" ) -const acquireAgentRuntimeLock = `-- name: AcquireAgentRuntimeLock :one -WITH lease_clock AS MATERIALIZED ( - SELECT statement_timestamp() AS renewed_at -) -INSERT INTO agent_runtime_locks(agent_id, worker_process_id, started_at, renewed_at, lease_expires_at) -SELECT agent.id, - $1::uuid, - lease_clock.renewed_at, - lease_clock.renewed_at, - lease_clock.renewed_at + ($2::bigint * interval '1 microsecond') -FROM agents agent -CROSS JOIN lease_clock -WHERE agent.project_id = $3 - AND agent.id = $4 - AND agent.state <> 'archived' -ON CONFLICT (agent_id) DO NOTHING -RETURNING id, agent_id, worker_process_id, started_at, renewed_at, - lease_expires_at, cancel_requested_at -` - -type AcquireAgentRuntimeLockParams struct { - WorkerProcessID uuid.UUID - LeaseDurationMicroseconds int64 - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) AcquireAgentRuntimeLock(ctx context.Context, arg AcquireAgentRuntimeLockParams) (AgentRuntimeLock, error) { - row := q.db.QueryRow(ctx, acquireAgentRuntimeLock, - arg.WorkerProcessID, - arg.LeaseDurationMicroseconds, - arg.ProjectID, - arg.AgentID, - ) - var i AgentRuntimeLock - err := row.Scan( - &i.ID, - &i.AgentID, - &i.WorkerProcessID, - &i.StartedAt, - &i.RenewedAt, - &i.LeaseExpiresAt, - &i.CancelRequestedAt, - ) - return i, err -} - -const agentRuntimeLockIsActive = `-- name: AgentRuntimeLockIsActive :one -SELECT EXISTS ( - SELECT 1 - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - AND agent.state <> 'archived' - WHERE agent.project_id = $1 - AND runtime_lock.agent_id = $2 - AND runtime_lock.id = $3 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) AS active -` - -type AgentRuntimeLockIsActiveParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) AgentRuntimeLockIsActive(ctx context.Context, arg AgentRuntimeLockIsActiveParams) (bool, error) { - row := q.db.QueryRow(ctx, agentRuntimeLockIsActive, arg.ProjectID, arg.AgentID, arg.ID) - var active bool - err := row.Scan(&active) - return active, err -} - const dBNow = `-- name: DBNow :one SELECT transaction_timestamp()::timestamptz ` @@ -98,29 +23,6 @@ func (q *Queries) DBNow(ctx context.Context) (time.Time, error) { return column_1, err } -const deleteAgentRuntimeLockForReap = `-- name: DeleteAgentRuntimeLockForReap :execrows -DELETE FROM agent_runtime_locks -USING agents agent -WHERE agent.project_id = $1 - AND agent.id = agent_runtime_locks.agent_id - AND agent_runtime_locks.agent_id = $2 - AND agent_runtime_locks.id = $3 -` - -type DeleteAgentRuntimeLockForReapParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) DeleteAgentRuntimeLockForReap(ctx context.Context, arg DeleteAgentRuntimeLockForReapParams) (int64, error) { - result, err := q.db.Exec(ctx, deleteAgentRuntimeLockForReap, arg.ProjectID, arg.AgentID, arg.ID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - const failQueuedProcessActionsForRuntimeEnd = `-- name: FailQueuedProcessActionsForRuntimeEnd :many UPDATE process_actions action SET state = 'failed', @@ -264,88 +166,6 @@ func (q *Queries) FailQueuedProcessesForRuntimeEnd(ctx context.Context, arg Fail return items, nil } -const failRuntimeToolCalls = `-- name: FailRuntimeToolCalls :many -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM tool_call_read_projection projection -WHERE call.agent_id = $1 - AND call.runtime_lock_id = $2::uuid - AND call.state = 'running' - AND call.type IN ('built_in', 'mcp') - AND projection.project_id = $3 - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, - projection.source_event_id, projection.model_call_context_id, call.provider_call_id, - call.name, call.input, - call.type, call.state, - 'failed'::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type FailRuntimeToolCallsParams struct { - AgentID uuid.UUID - RuntimeLockID uuid.UUID - ProjectID uuid.UUID -} - -type FailRuntimeToolCallsRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) FailRuntimeToolCalls(ctx context.Context, arg FailRuntimeToolCallsParams) ([]FailRuntimeToolCallsRow, error) { - rows, err := q.db.Query(ctx, failRuntimeToolCalls, arg.AgentID, arg.RuntimeLockID, arg.ProjectID) - if err != nil { - return nil, err - } - defer rows.Close() - items := []FailRuntimeToolCallsRow{} - for rows.Next() { - var i FailRuntimeToolCallsRow - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const getAgentRuntimeLockForRelease = `-- name: GetAgentRuntimeLockForRelease :one SELECT runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, runtime_lock.started_at, runtime_lock.renewed_at, @@ -378,37 +198,6 @@ func (q *Queries) GetAgentRuntimeLockForRelease(ctx context.Context, arg GetAgen return i, err } -const getPendingAgentRuntimeCancel = `-- name: GetPendingAgentRuntimeCancel :one -SELECT runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, - runtime_lock.started_at, runtime_lock.renewed_at, - runtime_lock.lease_expires_at, runtime_lock.cancel_requested_at -FROM agent_runtime_locks runtime_lock -JOIN agents agent ON agent.id = runtime_lock.agent_id -WHERE agent.project_id = $1 - AND runtime_lock.agent_id = $2 - AND runtime_lock.cancel_requested_at IS NOT NULL -` - -type GetPendingAgentRuntimeCancelParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) GetPendingAgentRuntimeCancel(ctx context.Context, arg GetPendingAgentRuntimeCancelParams) (AgentRuntimeLock, error) { - row := q.db.QueryRow(ctx, getPendingAgentRuntimeCancel, arg.ProjectID, arg.AgentID) - var i AgentRuntimeLock - err := row.Scan( - &i.ID, - &i.AgentID, - &i.WorkerProcessID, - &i.StartedAt, - &i.RenewedAt, - &i.LeaseExpiresAt, - &i.CancelRequestedAt, - ) - return i, err -} - const listExpiredAgentRuntimeLockCandidates = `-- name: ListExpiredAgentRuntimeLockCandidates :many SELECT runtime_lock.id, agent.project_id, runtime_lock.agent_id FROM agent_runtime_locks runtime_lock @@ -428,8 +217,6 @@ type ListExpiredAgentRuntimeLockCandidatesRow struct { AgentID uuid.UUID } -// Candidate discovery intentionally does not lock runtime rows. Destructive -// recovery must lock the canonical agent row before the runtime row. func (q *Queries) ListExpiredAgentRuntimeLockCandidates(ctx context.Context, arg ListExpiredAgentRuntimeLockCandidatesParams) ([]ListExpiredAgentRuntimeLockCandidatesRow, error) { rows, err := q.db.Query(ctx, listExpiredAgentRuntimeLockCandidates, arg.BatchSize) if err != nil { @@ -472,158 +259,3 @@ func (q *Queries) LockAgentRuntimeLockForOwnedMutation(ctx context.Context, arg err := row.Scan(&id) return id, err } - -const lockAgentRuntimeLockForRenewal = `-- name: LockAgentRuntimeLockForRenewal :one -SELECT runtime_lock.id -FROM agent_runtime_locks runtime_lock -JOIN agents agent ON agent.id = runtime_lock.agent_id -WHERE agent.project_id = $1 - AND runtime_lock.agent_id = $2 - AND runtime_lock.id = $3 -FOR UPDATE OF runtime_lock -` - -type LockAgentRuntimeLockForRenewalParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) LockAgentRuntimeLockForRenewal(ctx context.Context, arg LockAgentRuntimeLockForRenewalParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, lockAgentRuntimeLockForRenewal, arg.ProjectID, arg.AgentID, arg.ID) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const lockExpiredAgentRuntimeLockForReap = `-- name: LockExpiredAgentRuntimeLockForReap :one -SELECT runtime_lock.id, agent.project_id, runtime_lock.agent_id -FROM agent_runtime_locks runtime_lock -JOIN agents agent ON agent.id = runtime_lock.agent_id -WHERE agent.project_id = $1 - AND runtime_lock.agent_id = $2 - AND runtime_lock.id = $3 - AND runtime_lock.lease_expires_at <= statement_timestamp() -FOR UPDATE OF runtime_lock SKIP LOCKED -` - -type LockExpiredAgentRuntimeLockForReapParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -type LockExpiredAgentRuntimeLockForReapRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) LockExpiredAgentRuntimeLockForReap(ctx context.Context, arg LockExpiredAgentRuntimeLockForReapParams) (LockExpiredAgentRuntimeLockForReapRow, error) { - row := q.db.QueryRow(ctx, lockExpiredAgentRuntimeLockForReap, arg.ProjectID, arg.AgentID, arg.ID) - var i LockExpiredAgentRuntimeLockForReapRow - err := row.Scan(&i.ID, &i.ProjectID, &i.AgentID) - return i, err -} - -const releaseAgentRuntimeLock = `-- name: ReleaseAgentRuntimeLock :execrows -DELETE FROM agent_runtime_locks -USING agents agent -WHERE agent.project_id = $1 - AND agent.id = agent_runtime_locks.agent_id - AND agent_runtime_locks.agent_id = $2 - AND agent_runtime_locks.id = $3 -` - -type ReleaseAgentRuntimeLockParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) ReleaseAgentRuntimeLock(ctx context.Context, arg ReleaseAgentRuntimeLockParams) (int64, error) { - result, err := q.db.Exec(ctx, releaseAgentRuntimeLock, arg.ProjectID, arg.AgentID, arg.ID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const renewAgentRuntimeLock = `-- name: RenewAgentRuntimeLock :one -UPDATE agent_runtime_locks runtime_lock -SET renewed_at = greatest(runtime_lock.renewed_at, lease_clock.statement_at), - lease_expires_at = greatest( - runtime_lock.lease_expires_at, - greatest(runtime_lock.renewed_at, lease_clock.statement_at) + - ($1::bigint * interval '1 microsecond') - ) -FROM (SELECT statement_timestamp() AS statement_at) AS lease_clock -CROSS JOIN agents agent -WHERE agent.project_id = $2 - AND agent.id = runtime_lock.agent_id - AND runtime_lock.agent_id = $3 - AND runtime_lock.id = $4 -RETURNING runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, - runtime_lock.started_at, runtime_lock.renewed_at, - runtime_lock.lease_expires_at, runtime_lock.cancel_requested_at -` - -type RenewAgentRuntimeLockParams struct { - LeaseDurationMicroseconds int64 - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) RenewAgentRuntimeLock(ctx context.Context, arg RenewAgentRuntimeLockParams) (AgentRuntimeLock, error) { - row := q.db.QueryRow(ctx, renewAgentRuntimeLock, - arg.LeaseDurationMicroseconds, - arg.ProjectID, - arg.AgentID, - arg.ID, - ) - var i AgentRuntimeLock - err := row.Scan( - &i.ID, - &i.AgentID, - &i.WorkerProcessID, - &i.StartedAt, - &i.RenewedAt, - &i.LeaseExpiresAt, - &i.CancelRequestedAt, - ) - return i, err -} - -const requestAgentRuntimeCancel = `-- name: RequestAgentRuntimeCancel :one -UPDATE agent_runtime_locks runtime_lock -SET cancel_requested_at = coalesce(runtime_lock.cancel_requested_at, greatest(runtime_lock.started_at, statement_timestamp())) -FROM agents agent -WHERE agent.project_id = $1 - AND agent.id = runtime_lock.agent_id - AND runtime_lock.agent_id = $2 - AND runtime_lock.cancel_requested_at IS NULL -RETURNING runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, - runtime_lock.started_at, runtime_lock.renewed_at, - runtime_lock.lease_expires_at, runtime_lock.cancel_requested_at -` - -type RequestAgentRuntimeCancelParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) RequestAgentRuntimeCancel(ctx context.Context, arg RequestAgentRuntimeCancelParams) (AgentRuntimeLock, error) { - row := q.db.QueryRow(ctx, requestAgentRuntimeCancel, arg.ProjectID, arg.AgentID) - var i AgentRuntimeLock - err := row.Scan( - &i.ID, - &i.AgentID, - &i.WorkerProcessID, - &i.StartedAt, - &i.RenewedAt, - &i.LeaseExpiresAt, - &i.CancelRequestedAt, - ) - return i, err -} diff --git a/internal/storage/internal/dbsqlc/agent_turns.sql.go b/internal/storage/internal/dbsqlc/agent_turns.sql.go index 37b4f8b24..abcea0852 100644 --- a/internal/storage/internal/dbsqlc/agent_turns.sql.go +++ b/internal/storage/internal/dbsqlc/agent_turns.sql.go @@ -37,113 +37,30 @@ func (q *Queries) AgentTurnExistsInProject(ctx context.Context, arg AgentTurnExi return column_1, err } -const currentContinuableAgentTurn = `-- name: CurrentContinuableAgentTurn :one -WITH latest_turn AS MATERIALIZED ( - SELECT turn.id, agent.project_id, turn.agent_id, turn.turn_sequence, turn.latest_event_id, - turn.latest_semantic_event_id - FROM agent_turns turn - JOIN agents agent ON agent.id = turn.agent_id - WHERE agent.project_id = $1 - AND turn.agent_id = $2 - ORDER BY turn.turn_sequence DESC - LIMIT 1 -) -SELECT turn.id, turn.project_id, turn.agent_id, turn.turn_sequence, turn.latest_event_id, - turn.latest_semantic_event_id -FROM latest_turn turn -WHERE EXISTS ( - SELECT 1 - FROM agent_continuable_model_contexts(turn.project_id, turn.agent_id) - AS context(turn_id, model_call_context_id, input_event_sequence, has_later_semantic_event) - WHERE context.turn_id = turn.id - ) - OR agent_has_incomplete_tool_batch(turn.project_id, turn.agent_id) - OR EXISTS ( - SELECT 1 - FROM agent_next_model_work(turn.project_id, turn.agent_id) frontier - WHERE frontier.turn_id = turn.id - ) +const getAgentTurn = `-- name: GetAgentTurn :one +SELECT t.id,t.turn_sequence,t.latest_event_id,t.latest_semantic_event_id +FROM agent_turns t JOIN agents a ON a.id=t.agent_id +WHERE a.project_id=$1 AND t.agent_id=$2 AND t.id=$3 ` -type CurrentContinuableAgentTurnParams struct { +type GetAgentTurnParams struct { ProjectID uuid.UUID AgentID uuid.UUID + ID uuid.UUID } -type CurrentContinuableAgentTurnRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnSequence int64 - LatestEventID uuid.UUID - LatestSemanticEventID uuid.UUID -} - -func (q *Queries) CurrentContinuableAgentTurn(ctx context.Context, arg CurrentContinuableAgentTurnParams) (CurrentContinuableAgentTurnRow, error) { - row := q.db.QueryRow(ctx, currentContinuableAgentTurn, arg.ProjectID, arg.AgentID) - var i CurrentContinuableAgentTurnRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnSequence, - &i.LatestEventID, - &i.LatestSemanticEventID, - ) - return i, err -} - -const insertAgentTurn = `-- name: InsertAgentTurn :one -WITH inserted AS ( - INSERT INTO agent_turns(id, agent_id, turn_sequence, latest_event_id, latest_semantic_event_id) - SELECT $1, $2, - $3, event.id, $4 - FROM agent_events event - JOIN agents agent ON agent.id = event.agent_id - WHERE agent.project_id = $5 - AND event.agent_id = $2 - AND event.turn_id = $1 - AND event.id = $6 - RETURNING id, agent_id, turn_sequence, latest_event_id, latest_semantic_event_id -) -SELECT inserted.id, agent.project_id, inserted.agent_id, inserted.turn_sequence, - inserted.latest_event_id, inserted.latest_semantic_event_id -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id -` - -type InsertAgentTurnParams struct { - ID uuid.UUID - AgentID uuid.UUID - TurnSequence int64 - LatestSemanticEventID uuid.UUID - ProjectID uuid.UUID - LatestEventID uuid.UUID -} - -type InsertAgentTurnRow struct { +type GetAgentTurnRow struct { ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID TurnSequence int64 LatestEventID uuid.UUID LatestSemanticEventID uuid.UUID } -func (q *Queries) InsertAgentTurn(ctx context.Context, arg InsertAgentTurnParams) (InsertAgentTurnRow, error) { - row := q.db.QueryRow(ctx, insertAgentTurn, - arg.ID, - arg.AgentID, - arg.TurnSequence, - arg.LatestSemanticEventID, - arg.ProjectID, - arg.LatestEventID, - ) - var i InsertAgentTurnRow +func (q *Queries) GetAgentTurn(ctx context.Context, arg GetAgentTurnParams) (GetAgentTurnRow, error) { + row := q.db.QueryRow(ctx, getAgentTurn, arg.ProjectID, arg.AgentID, arg.ID) + var i GetAgentTurnRow err := row.Scan( &i.ID, - &i.ProjectID, - &i.AgentID, &i.TurnSequence, &i.LatestEventID, &i.LatestSemanticEventID, @@ -764,63 +681,3 @@ func (q *Queries) ListTurnEventsForRead(ctx context.Context, arg ListTurnEventsF } return items, nil } - -const nextTurnSequence = `-- name: NextTurnSequence :one -SELECT (coalesce(max(turn.turn_sequence), 0) + 1)::bigint -FROM agent_turns turn -JOIN agents agent ON agent.id = turn.agent_id -WHERE agent.project_id = $1 - AND turn.agent_id = $2 -` - -type NextTurnSequenceParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) NextTurnSequence(ctx context.Context, arg NextTurnSequenceParams) (int64, error) { - row := q.db.QueryRow(ctx, nextTurnSequence, arg.ProjectID, arg.AgentID) - var column_1 int64 - err := row.Scan(&column_1) - return column_1, err -} - -const updateAgentTurnLatestEvent = `-- name: UpdateAgentTurnLatestEvent :execrows -UPDATE agent_turns -SET latest_event_id = $1, - latest_semantic_event_id = coalesce($2, latest_semantic_event_id) -FROM agent_events event -WHERE agent_turns.agent_id = $3 - AND agent_turns.id = $4 - AND event.agent_id = agent_turns.agent_id - AND event.turn_id = agent_turns.id - AND event.id = $1 - AND EXISTS ( - SELECT 1 - FROM agents agent - WHERE agent.id = agent_turns.agent_id - AND agent.project_id = $5 - ) -` - -type UpdateAgentTurnLatestEventParams struct { - LatestEventID uuid.UUID - LatestSemanticEventID *uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - ProjectID uuid.UUID -} - -func (q *Queries) UpdateAgentTurnLatestEvent(ctx context.Context, arg UpdateAgentTurnLatestEventParams) (int64, error) { - result, err := q.db.Exec(ctx, updateAgentTurnLatestEvent, - arg.LatestEventID, - arg.LatestSemanticEventID, - arg.AgentID, - arg.ID, - arg.ProjectID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} diff --git a/internal/storage/internal/dbsqlc/agent_work.sql.go b/internal/storage/internal/dbsqlc/agent_work.sql.go deleted file mode 100644 index ffc6f6367..000000000 --- a/internal/storage/internal/dbsqlc/agent_work.sql.go +++ /dev/null @@ -1,131 +0,0 @@ -// Code generated by sqlc. DO NOT EDIT. -// versions: -// sqlc v1.31.1 -// source: agent_work.sql - -package dbsqlc - -import ( - "context" - - "github.com/google/uuid" -) - -const agentHasIncompleteToolBatch = `-- name: AgentHasIncompleteToolBatch :one -SELECT agent_has_incomplete_tool_batch( - $1, - $2 -) AS incomplete -` - -type AgentHasIncompleteToolBatchParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) AgentHasIncompleteToolBatch(ctx context.Context, arg AgentHasIncompleteToolBatchParams) (bool, error) { - row := q.db.QueryRow(ctx, agentHasIncompleteToolBatch, arg.ProjectID, arg.AgentID) - var incomplete bool - err := row.Scan(&incomplete) - return incomplete, err -} - -const nextAgentModelWork = `-- name: NextAgentModelWork :one -SELECT frontier.work_kind::text AS work_kind, - coalesce(frontier.model_call_context_id, '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS model_call_context_id, - coalesce(frontier.model_output_id, '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS model_output_id, - frontier.turn_id::uuid AS turn_id, - frontier.input_ids::uuid[] AS input_ids, - frontier.opening_event_sequence::bigint AS opening_event_sequence, - frontier.ready_at <= statement_timestamp() AS is_ready -FROM agent_next_model_work( - $1, - $2 -) AS frontier( - work_kind, - model_call_context_id, - model_output_id, - turn_id, - input_ids, - opening_event_sequence, - ready_at -) -LIMIT 1 -` - -type NextAgentModelWorkParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type NextAgentModelWorkRow struct { - WorkKind string - ModelCallContextID uuid.UUID - ModelOutputID uuid.UUID - TurnID uuid.UUID - InputIds []uuid.UUID - OpeningEventSequence int64 - IsReady bool -} - -func (q *Queries) NextAgentModelWork(ctx context.Context, arg NextAgentModelWorkParams) (NextAgentModelWorkRow, error) { - row := q.db.QueryRow(ctx, nextAgentModelWork, arg.ProjectID, arg.AgentID) - var i NextAgentModelWorkRow - err := row.Scan( - &i.WorkKind, - &i.ModelCallContextID, - &i.ModelOutputID, - &i.TurnID, - &i.InputIds, - &i.OpeningEventSequence, - &i.IsReady, - ) - return i, err -} - -const nextAgentToolWork = `-- name: NextAgentToolWork :one -SELECT frontier.turn_id::uuid AS turn_id, - frontier.model_call_context_id::uuid AS model_call_context_id, - frontier.model_output_id::uuid AS model_output_id, - frontier.source_event_id::uuid AS source_event_id -FROM agent_tool_work_frontiers( - $1, - $2 -) AS frontier( - turn_id, - model_call_context_id, - model_output_id, - source_event_id, - ready_at, - frontier_order_created_at, - frontier_order_tool_call_id -) -ORDER BY frontier.ready_at, - frontier.frontier_order_created_at, - frontier.frontier_order_tool_call_id -LIMIT 1 -` - -type NextAgentToolWorkParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type NextAgentToolWorkRow struct { - TurnID uuid.UUID - ModelCallContextID uuid.UUID - ModelOutputID uuid.UUID - SourceEventID uuid.UUID -} - -func (q *Queries) NextAgentToolWork(ctx context.Context, arg NextAgentToolWorkParams) (NextAgentToolWorkRow, error) { - row := q.db.QueryRow(ctx, nextAgentToolWork, arg.ProjectID, arg.AgentID) - var i NextAgentToolWorkRow - err := row.Scan( - &i.TurnID, - &i.ModelCallContextID, - &i.ModelOutputID, - &i.SourceEventID, - ) - return i, err -} diff --git a/internal/storage/internal/dbsqlc/agents.sql.go b/internal/storage/internal/dbsqlc/agents.sql.go index 84cb0d0d0..c27d985ef 100644 --- a/internal/storage/internal/dbsqlc/agents.sql.go +++ b/internal/storage/internal/dbsqlc/agents.sql.go @@ -1193,28 +1193,6 @@ func (q *Queries) RetargetAgentProfile(ctx context.Context, arg RetargetAgentPro return i, err } -const updateAgentCurrentConfig = `-- name: UpdateAgentCurrentConfig :execrows -UPDATE agents -SET current_config_id = $1, - updated_at = statement_timestamp() -WHERE project_id = $2 - AND id = $3 -` - -type UpdateAgentCurrentConfigParams struct { - AgentConfigID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) UpdateAgentCurrentConfig(ctx context.Context, arg UpdateAgentCurrentConfigParams) (int64, error) { - result, err := q.db.Exec(ctx, updateAgentCurrentConfig, arg.AgentConfigID, arg.ProjectID, arg.AgentID) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - const upsertAgentConfigByHash = `-- name: UpsertAgentConfigByHash :one WITH inserted_config AS ( INSERT INTO agent_configs( diff --git a/internal/storage/internal/dbsqlc/context_checkpoints.sql.go b/internal/storage/internal/dbsqlc/context_checkpoints.sql.go index 7d1fe7729..61b77373a 100644 --- a/internal/storage/internal/dbsqlc/context_checkpoints.sql.go +++ b/internal/storage/internal/dbsqlc/context_checkpoints.sql.go @@ -12,70 +12,6 @@ import ( "github.com/google/uuid" ) -const countCheckpointRangeEvents = `-- name: CountCheckpointRangeEvents :one -SELECT count(*)::bigint AS count -FROM agent_events -JOIN agents agent ON agent.id = agent_events.agent_id -WHERE agent.project_id = $1 - AND agent_events.agent_id = $2 - AND sequence BETWEEN $3 AND $4 -` - -type CountCheckpointRangeEventsParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - StartSequence int64 - EndSequence int64 -} - -func (q *Queries) CountCheckpointRangeEvents(ctx context.Context, arg CountCheckpointRangeEventsParams) (int64, error) { - row := q.db.QueryRow(ctx, countCheckpointRangeEvents, - arg.ProjectID, - arg.AgentID, - arg.StartSequence, - arg.EndSequence, - ) - var count int64 - err := row.Scan(&count) - return count, err -} - -const countCheckpointRangeOpenToolResults = `-- name: CountCheckpointRangeOpenToolResults :one -SELECT count(*)::bigint AS count -FROM tool_call_read_projection tool_call -LEFT JOIN tool_call_results result ON result.agent_id = tool_call.agent_id - AND result.tool_call_id = tool_call.id -LEFT JOIN agent_events result_event ON result_event.agent_id = result.agent_id - AND result_event.tool_call_result_id = result.id - AND result_event.event_kind = 'tool_result' -WHERE tool_call.project_id = $1 - AND tool_call.agent_id = $2 - AND tool_call.source_event_sequence BETWEEN $3 AND $4 - AND ( - result_event.id IS NULL - OR result_event.sequence > $4 - ) -` - -type CountCheckpointRangeOpenToolResultsParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - StartSequence int64 - EndSequence int64 -} - -func (q *Queries) CountCheckpointRangeOpenToolResults(ctx context.Context, arg CountCheckpointRangeOpenToolResultsParams) (int64, error) { - row := q.db.QueryRow(ctx, countCheckpointRangeOpenToolResults, - arg.ProjectID, - arg.AgentID, - arg.StartSequence, - arg.EndSequence, - ) - var count int64 - err := row.Scan(&count) - return count, err -} - const countConsecutiveContextCheckpointLineage = `-- name: CountConsecutiveContextCheckpointLineage :one WITH RECURSIVE lineage AS ( SELECT producer.input_event_sequence AS prior_frontier @@ -277,49 +213,3 @@ func (q *Queries) GetLatestApplicableContextCheckpoint(ctx context.Context, arg ) return i, err } - -const insertContextCheckpoint = `-- name: InsertContextCheckpoint :one -INSERT INTO context_checkpoints( - id, agent_id, summarized_through_event_sequence, - producer_model_call_context_id, summary, created_at -) -SELECT $1, agent.id, - context.source_event_sequence_end, - context.id, - $2, - statement_timestamp() -FROM agents agent -JOIN model_call_contexts context ON context.project_id = agent.project_id - AND context.agent_id = agent.id - AND context.id = $3 - AND context.operation_kind = 'compaction' - AND context.state = 'started' -JOIN agent_runtime_locks runtime_lock ON runtime_lock.agent_id = context.agent_id - AND runtime_lock.id = context.runtime_lock_id - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -WHERE agent.project_id = $4 - AND agent.id = $5 -RETURNING id -` - -type InsertContextCheckpointParams struct { - ID uuid.UUID - Summary string - ProducerModelCallContextID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) InsertContextCheckpoint(ctx context.Context, arg InsertContextCheckpointParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, insertContextCheckpoint, - arg.ID, - arg.Summary, - arg.ProducerModelCallContextID, - arg.ProjectID, - arg.AgentID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} diff --git a/internal/storage/internal/dbsqlc/events.sql.go b/internal/storage/internal/dbsqlc/events.sql.go index 6068a29b5..a83c6a136 100644 --- a/internal/storage/internal/dbsqlc/events.sql.go +++ b/internal/storage/internal/dbsqlc/events.sql.go @@ -13,46 +13,6 @@ import ( "github.com/google/uuid" ) -const advanceEventSequence = `-- name: AdvanceEventSequence :exec -UPDATE agents -SET next_event_sequence = next_event_sequence + 1, updated_at = statement_timestamp() -WHERE project_id = $1 - AND id = $2 -` - -type AdvanceEventSequenceParams struct { - ProjectID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) AdvanceEventSequence(ctx context.Context, arg AdvanceEventSequenceParams) error { - _, err := q.db.Exec(ctx, advanceEventSequence, arg.ProjectID, arg.ID) - return err -} - -const allocateEventSequence = `-- name: AllocateEventSequence :one -SELECT project_id, next_event_sequence -FROM agents AS agents -WHERE id = $1 -FOR UPDATE -` - -type AllocateEventSequenceParams struct { - ID uuid.UUID -} - -type AllocateEventSequenceRow struct { - ProjectID uuid.UUID - NextEventSequence int64 -} - -func (q *Queries) AllocateEventSequence(ctx context.Context, arg AllocateEventSequenceParams) (AllocateEventSequenceRow, error) { - row := q.db.QueryRow(ctx, allocateEventSequence, arg.ID) - var i AllocateEventSequenceRow - err := row.Scan(&i.ProjectID, &i.NextEventSequence) - return i, err -} - const getEventByProjectAgentIdempotencyKey = `-- name: GetEventByProjectAgentIdempotencyKey :one SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, event.sequence, event.event_kind, event.created_at, coalesce(event.idempotency_key, '') AS idempotency_key FROM agent_events event @@ -93,47 +53,6 @@ func (q *Queries) GetEventByProjectAgentIdempotencyKey(ctx context.Context, arg return i, err } -const latestAgentEvent = `-- name: LatestAgentEvent :one -SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, event.sequence, event.event_kind, event.created_at, coalesce(event.idempotency_key, '') AS idempotency_key -FROM agent_events event -JOIN agents agent ON agent.id = event.agent_id -WHERE agent.project_id = $1 AND event.agent_id = $2 -ORDER BY sequence DESC -LIMIT 1 -` - -type LatestAgentEventParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type LatestAgentEventRow struct { - ID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - IsOpeningEvent bool - Sequence int64 - EventKind string - CreatedAt time.Time - IdempotencyKey string -} - -func (q *Queries) LatestAgentEvent(ctx context.Context, arg LatestAgentEventParams) (LatestAgentEventRow, error) { - row := q.db.QueryRow(ctx, latestAgentEvent, arg.ProjectID, arg.AgentID) - var i LatestAgentEventRow - err := row.Scan( - &i.ID, - &i.AgentID, - &i.TurnID, - &i.IsOpeningEvent, - &i.Sequence, - &i.EventKind, - &i.CreatedAt, - &i.IdempotencyKey, - ) - return i, err -} - const listAgentEventFrontiers = `-- name: ListAgentEventFrontiers :many SELECT id AS agent_id, (next_event_sequence - 1)::bigint AS event_sequence FROM agents diff --git a/internal/storage/internal/dbsqlc/integration_interactions.sql.go b/internal/storage/internal/dbsqlc/integration_interactions.sql.go index 7013c1542..6e86ac927 100644 --- a/internal/storage/internal/dbsqlc/integration_interactions.sql.go +++ b/internal/storage/internal/dbsqlc/integration_interactions.sql.go @@ -7,7 +7,6 @@ package dbsqlc import ( "context" - "encoding/json" "github.com/google/uuid" ) @@ -120,78 +119,3 @@ func (q *Queries) GetInteractionSelection(ctx context.Context, arg GetInteractio ) return i, err } - -const recordAgentInteractionPresentationReceipt = `-- name: RecordAgentInteractionPresentationReceipt :execrows -UPDATE agent_interactions interaction -SET presentation_receipt = $1::jsonb -WHERE interaction.agent_id = $2 AND interaction.id = $3 - AND interaction.destination = $4::jsonb - AND interaction.presentation_receipt IS NULL - AND EXISTS ( - SELECT 1 FROM agents agent - WHERE agent.project_id = $5 AND agent.id = interaction.agent_id - ) -` - -type RecordAgentInteractionPresentationReceiptParams struct { - Receipt json.RawMessage - AgentID uuid.UUID - ID uuid.UUID - Destination json.RawMessage - ProjectID uuid.UUID -} - -func (q *Queries) RecordAgentInteractionPresentationReceipt(ctx context.Context, arg RecordAgentInteractionPresentationReceiptParams) (int64, error) { - result, err := q.db.Exec(ctx, recordAgentInteractionPresentationReceipt, - arg.Receipt, - arg.AgentID, - arg.ID, - arg.Destination, - arg.ProjectID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const setInteractionSelection = `-- name: SetInteractionSelection :execrows -UPDATE agents -SET interaction_target_id = $1::uuid, - interaction_handler_key = $2::text, - interaction_auto_select = $3::boolean, - updated_at = statement_timestamp() -WHERE agents.project_id = $4 AND agents.id = $5 - AND (($1::uuid IS NULL AND $2::text IS NULL) - OR ($2::text <> '' AND EXISTS ( - SELECT 1 FROM integration_targets target - JOIN integrations integration - ON integration.project_id = target.project_id - AND integration.id = target.integration_id - WHERE target.project_id = agents.project_id AND target.agent_id = agents.id - AND target.id = $1::uuid AND target.deleted_at IS NULL - AND integration.deleted_at IS NULL AND integration.state = 'active' - ))) -` - -type SetInteractionSelectionParams struct { - TargetID *uuid.UUID - HandlerKey *string - AutoSelect bool - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) SetInteractionSelection(ctx context.Context, arg SetInteractionSelectionParams) (int64, error) { - result, err := q.db.Exec(ctx, setInteractionSelection, - arg.TargetID, - arg.HandlerKey, - arg.AutoSelect, - arg.ProjectID, - arg.AgentID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} diff --git a/internal/storage/internal/dbsqlc/integration_target_lifecycle.sql.go b/internal/storage/internal/dbsqlc/integration_target_lifecycle.sql.go index 66df3005a..06413e306 100644 --- a/internal/storage/internal/dbsqlc/integration_target_lifecycle.sql.go +++ b/internal/storage/internal/dbsqlc/integration_target_lifecycle.sql.go @@ -12,27 +12,6 @@ import ( "github.com/google/uuid" ) -const clearDeletedIntegrationTargetsFromAgents = `-- name: ClearDeletedIntegrationTargetsFromAgents :exec -UPDATE agents agent SET interaction_target_id = NULL, interaction_handler_key = NULL, updated_at = statement_timestamp() -FROM integration_targets target -WHERE agent.project_id = $1 - AND target.project_id = agent.project_id - AND target.integration_id = $2 - AND agent.id = target.agent_id - AND agent.interaction_target_id = target.id -` - -type ClearDeletedIntegrationTargetsFromAgentsParams struct { - ProjectID uuid.UUID - IntegrationID uuid.UUID -} - -// @sqlc-vet-disable integration-targets-deleted-at -func (q *Queries) ClearDeletedIntegrationTargetsFromAgents(ctx context.Context, arg ClearDeletedIntegrationTargetsFromAgentsParams) error { - _, err := q.db.Exec(ctx, clearDeletedIntegrationTargetsFromAgents, arg.ProjectID, arg.IntegrationID) - return err -} - const deleteIntegrationTargets = `-- name: DeleteIntegrationTargets :exec UPDATE integration_targets SET deleted_at = statement_timestamp(), updated_at = statement_timestamp() WHERE project_id = $1 AND integration_id = $2 diff --git a/internal/storage/internal/dbsqlc/kernel_event_authority.sql.go b/internal/storage/internal/dbsqlc/kernel_event_authority.sql.go index 05f598f54..eade02ad2 100644 --- a/internal/storage/internal/dbsqlc/kernel_event_authority.sql.go +++ b/internal/storage/internal/dbsqlc/kernel_event_authority.sql.go @@ -15,13 +15,7 @@ import ( const getModelOutputByModelContext = `-- name: GetModelOutputByModelContext :one SELECT output.id, agent.project_id, output.agent_id, - ( - SELECT context_turn.turn_id - FROM model_call_context_turns context_turn - WHERE context_turn.project_id = agent.project_id - AND context_turn.agent_id = output.agent_id - AND context_turn.model_call_context_id = output.model_call_context_id - ) AS turn_id, + context.turn_id, output.model_call_context_id, output.served_provider_model_slug, output.stop_reason, context.provider_response_id, @@ -89,27 +83,6 @@ func (q *Queries) GetModelOutputByModelContext(ctx context.Context, arg GetModel return i, err } -const getToolCallProviderCallID = `-- name: GetToolCallProviderCallID :one -SELECT provider_call_id -FROM tool_call_read_projection -WHERE project_id = $1 - AND agent_id = $2 - AND id = $3 -` - -type GetToolCallProviderCallIDParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) GetToolCallProviderCallID(ctx context.Context, arg GetToolCallProviderCallIDParams) (string, error) { - row := q.db.QueryRow(ctx, getToolCallProviderCallID, arg.ProjectID, arg.AgentID, arg.ID) - var provider_call_id string - err := row.Scan(&provider_call_id) - return provider_call_id, err -} - const getToolCallResultByToolCall = `-- name: GetToolCallResultByToolCall :one SELECT result.id, agent.project_id, result.agent_id, source_event.turn_id AS turn_id, @@ -157,60 +130,6 @@ func (q *Queries) GetToolCallResultByToolCall(ctx context.Context, arg GetToolCa return i, err } -const getTypedAgentEventByIdempotency = `-- name: GetTypedAgentEventByIdempotency :one -SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, - event.sequence, event.event_kind, event.created_at, - coalesce(event.idempotency_key, '') AS idempotency_key, - event.agent_input_id, event.model_output_id, event.tool_call_result_id, - event.context_checkpoint_id -FROM agent_events event -JOIN agents agent ON agent.id = event.agent_id -WHERE agent.project_id = $1 - AND event.agent_id = $2 - AND event.idempotency_key = $3::text -` - -type GetTypedAgentEventByIdempotencyParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - IdempotencyKey string -} - -type GetTypedAgentEventByIdempotencyRow struct { - ID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - IsOpeningEvent bool - Sequence int64 - EventKind string - CreatedAt time.Time - IdempotencyKey string - AgentInputID *uuid.UUID - ModelOutputID *uuid.UUID - ToolCallResultID *uuid.UUID - ContextCheckpointID *uuid.UUID -} - -func (q *Queries) GetTypedAgentEventByIdempotency(ctx context.Context, arg GetTypedAgentEventByIdempotencyParams) (GetTypedAgentEventByIdempotencyRow, error) { - row := q.db.QueryRow(ctx, getTypedAgentEventByIdempotency, arg.ProjectID, arg.AgentID, arg.IdempotencyKey) - var i GetTypedAgentEventByIdempotencyRow - err := row.Scan( - &i.ID, - &i.AgentID, - &i.TurnID, - &i.IsOpeningEvent, - &i.Sequence, - &i.EventKind, - &i.CreatedAt, - &i.IdempotencyKey, - &i.AgentInputID, - &i.ModelOutputID, - &i.ToolCallResultID, - &i.ContextCheckpointID, - ) - return i, err -} - const getTypedAgentEventByModelOutput = `-- name: GetTypedAgentEventByModelOutput :one SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, event.sequence, event.event_kind, event.created_at, @@ -265,629 +184,6 @@ func (q *Queries) GetTypedAgentEventByModelOutput(ctx context.Context, arg GetTy return i, err } -const insertConfigChangeAgentInput = `-- name: InsertConfigChangeAgentInput :one -WITH generated AS ( - SELECT uuidv7() AS id -), -inserted AS ( -INSERT INTO agent_inputs( - id, project_id, agent_id, state, - actor_id, input_kind, delivery_mode, - agent_config_id, idempotency_scope, input_idempotency_key, queued_at, metadata -) -SELECT generated.id, agent.project_id, agent.id, 'received', - $1::uuid, 'config_change', 'immediate', - $2, - $3, $4, - statement_timestamp(), $5 -FROM agents agent -JOIN generated ON true -WHERE agent.project_id = $6 - AND agent.id = $7 -ON CONFLICT DO NOTHING -RETURNING id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - coalesce(idempotency_scope, '') AS idempotency_scope, - coalesce(input_idempotency_key, '') AS input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, - coalesce(rejected_reason, '') AS rejected_reason, metadata -) -SELECT id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - idempotency_scope, - input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - control_type, target_interaction_id, agent_config_id, resolved_at, - rejected_reason, metadata -FROM inserted -UNION ALL -SELECT id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - coalesce(idempotency_scope, '') AS idempotency_scope, - coalesce(input_idempotency_key, '') AS input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, - coalesce(rejected_reason, '') AS rejected_reason, metadata -FROM agent_inputs -WHERE project_id = $6 - AND agent_id = $7 - AND idempotency_scope = $3 - AND input_idempotency_key = $4 -LIMIT 1 -` - -type InsertConfigChangeAgentInputParams struct { - ActorID *uuid.UUID - AgentConfigID *uuid.UUID - IdempotencyScope *string - InputIdempotencyKey *string - Metadata json.RawMessage - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type InsertConfigChangeAgentInputRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - State string - InputRank int64 - ActorID *uuid.UUID - InputKind string - IdempotencyScope string - InputIdempotencyKey string - QueuedAt time.Time - AdmittedEventID *uuid.UUID - AdmittedAt *time.Time - CanceledAt *time.Time - DeliveryMode string - ControlType string - TargetInteractionID *uuid.UUID - AgentConfigID *uuid.UUID - ResolvedAt *time.Time - RejectedReason string - Metadata json.RawMessage -} - -func (q *Queries) InsertConfigChangeAgentInput(ctx context.Context, arg InsertConfigChangeAgentInputParams) (InsertConfigChangeAgentInputRow, error) { - row := q.db.QueryRow(ctx, insertConfigChangeAgentInput, - arg.ActorID, - arg.AgentConfigID, - arg.IdempotencyScope, - arg.InputIdempotencyKey, - arg.Metadata, - arg.ProjectID, - arg.AgentID, - ) - var i InsertConfigChangeAgentInputRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.State, - &i.InputRank, - &i.ActorID, - &i.InputKind, - &i.IdempotencyScope, - &i.InputIdempotencyKey, - &i.QueuedAt, - &i.AdmittedEventID, - &i.AdmittedAt, - &i.CanceledAt, - &i.DeliveryMode, - &i.ControlType, - &i.TargetInteractionID, - &i.AgentConfigID, - &i.ResolvedAt, - &i.RejectedReason, - &i.Metadata, - ) - return i, err -} - -const insertContentBlock = `-- name: InsertContentBlock :one -WITH content_owner AS ( - SELECT input.queued_at AS created_at - FROM agent_inputs input - WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3::uuid - - UNION ALL - - SELECT output.created_at - FROM model_outputs output - JOIN agents agent ON agent.id = output.agent_id - WHERE agent.project_id = $1 - AND output.agent_id = $2 - AND output.id = $4::uuid - - UNION ALL - - SELECT result.completed_at - FROM tool_call_results result - JOIN agents agent ON agent.id = result.agent_id - WHERE agent.project_id = $1 - AND result.agent_id = $2 - AND result.id = $5::uuid -), -inserted AS ( - INSERT INTO content_blocks( - agent_id, owner_kind, owner_agent_input_id, - owner_model_output_id, owner_tool_call_result_id, ordinal, block_kind, - text_content, structured_data, artifact_id, tool_call_id, - exclude_from_model_context, metadata, created_at - ) - SELECT - $2, $6, - $3::uuid, $4::uuid, - $5::uuid, $7, $8, - $9, $10::jsonb, $11::uuid, $12::uuid, - $13, $14::jsonb, content_owner.created_at - FROM content_owner - RETURNING id, agent_id, owner_kind, owner_agent_input_id, - owner_model_output_id, owner_tool_call_result_id, ordinal, block_kind, - coalesce(text_content, '') AS text_content, structured_data, - artifact_id, tool_call_id, created_at -) -SELECT inserted.id, agent.project_id, inserted.agent_id, inserted.owner_kind, - inserted.owner_agent_input_id, inserted.owner_model_output_id, - inserted.owner_tool_call_result_id, inserted.ordinal, inserted.block_kind, - inserted.text_content, inserted.structured_data, inserted.artifact_id, - inserted.tool_call_id, inserted.created_at -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id -` - -type InsertContentBlockParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - OwnerAgentInputID *uuid.UUID - OwnerModelOutputID *uuid.UUID - OwnerToolCallResultID *uuid.UUID - OwnerKind string - Ordinal int32 - BlockKind string - TextContent *string - StructuredData *json.RawMessage - ArtifactID *uuid.UUID - ToolCallID *uuid.UUID - ExcludeFromModelContext bool - Metadata json.RawMessage -} - -type InsertContentBlockRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - OwnerKind string - OwnerAgentInputID *uuid.UUID - OwnerModelOutputID *uuid.UUID - OwnerToolCallResultID *uuid.UUID - Ordinal int32 - BlockKind string - TextContent string - StructuredData *json.RawMessage - ArtifactID *uuid.UUID - ToolCallID *uuid.UUID - CreatedAt time.Time -} - -func (q *Queries) InsertContentBlock(ctx context.Context, arg InsertContentBlockParams) (InsertContentBlockRow, error) { - row := q.db.QueryRow(ctx, insertContentBlock, - arg.ProjectID, - arg.AgentID, - arg.OwnerAgentInputID, - arg.OwnerModelOutputID, - arg.OwnerToolCallResultID, - arg.OwnerKind, - arg.Ordinal, - arg.BlockKind, - arg.TextContent, - arg.StructuredData, - arg.ArtifactID, - arg.ToolCallID, - arg.ExcludeFromModelContext, - arg.Metadata, - ) - var i InsertContentBlockRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.OwnerKind, - &i.OwnerAgentInputID, - &i.OwnerModelOutputID, - &i.OwnerToolCallResultID, - &i.Ordinal, - &i.BlockKind, - &i.TextContent, - &i.StructuredData, - &i.ArtifactID, - &i.ToolCallID, - &i.CreatedAt, - ) - return i, err -} - -const insertControlAgentInput = `-- name: InsertControlAgentInput :one -WITH generated AS ( - SELECT uuidv7() AS id -) -INSERT INTO agent_inputs( - id, project_id, agent_id, state, - actor_id, input_kind, delivery_mode, control_type, - idempotency_scope, input_idempotency_key, queued_at, metadata -) -SELECT generated.id, agent.project_id, agent.id, 'received', - $1, 'control', 'immediate', - $2, $3, $4, - statement_timestamp(), $5 -FROM agents agent -JOIN generated ON true -WHERE agent.project_id = $6 - AND agent.id = $7 -ON CONFLICT DO NOTHING -RETURNING id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - coalesce(idempotency_scope, '') AS idempotency_scope, - coalesce(input_idempotency_key, '') AS input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, - coalesce(rejected_reason, '') AS rejected_reason, metadata -` - -type InsertControlAgentInputParams struct { - ActorID *uuid.UUID - ControlType *string - IdempotencyScope *string - InputIdempotencyKey *string - Metadata json.RawMessage - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type InsertControlAgentInputRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - State string - InputRank int64 - ActorID *uuid.UUID - InputKind string - IdempotencyScope string - InputIdempotencyKey string - QueuedAt time.Time - AdmittedEventID *uuid.UUID - AdmittedAt *time.Time - CanceledAt *time.Time - DeliveryMode string - ControlType string - TargetInteractionID *uuid.UUID - AgentConfigID *uuid.UUID - ResolvedAt *time.Time - RejectedReason string - Metadata json.RawMessage -} - -func (q *Queries) InsertControlAgentInput(ctx context.Context, arg InsertControlAgentInputParams) (InsertControlAgentInputRow, error) { - row := q.db.QueryRow(ctx, insertControlAgentInput, - arg.ActorID, - arg.ControlType, - arg.IdempotencyScope, - arg.InputIdempotencyKey, - arg.Metadata, - arg.ProjectID, - arg.AgentID, - ) - var i InsertControlAgentInputRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.State, - &i.InputRank, - &i.ActorID, - &i.InputKind, - &i.IdempotencyScope, - &i.InputIdempotencyKey, - &i.QueuedAt, - &i.AdmittedEventID, - &i.AdmittedAt, - &i.CanceledAt, - &i.DeliveryMode, - &i.ControlType, - &i.TargetInteractionID, - &i.AgentConfigID, - &i.ResolvedAt, - &i.RejectedReason, - &i.Metadata, - ) - return i, err -} - -const insertModelOutputAuthority = `-- name: InsertModelOutputAuthority :one -WITH inserted AS ( - INSERT INTO model_outputs( - agent_id, model_call_context_id, - served_provider_model_slug, stop_reason, provider_replay, created_at - ) - SELECT agent.id, context.id, - $1, $2, - $3::jsonb, statement_timestamp() - FROM agents agent - JOIN model_call_contexts context ON context.project_id = agent.project_id - AND context.agent_id = agent.id - AND context.id = $4::uuid - WHERE agent.project_id = $5 - AND agent.id = $6 - ON CONFLICT (agent_id, model_call_context_id) DO NOTHING - RETURNING id, agent_id, model_call_context_id, - served_provider_model_slug, stop_reason, provider_replay, created_at -) -SELECT inserted.id, agent.project_id, inserted.agent_id, - ( - SELECT context_turn.turn_id - FROM model_call_context_turns context_turn - WHERE context_turn.project_id = agent.project_id - AND context_turn.agent_id = inserted.agent_id - AND context_turn.model_call_context_id = inserted.model_call_context_id - ) AS turn_id, - inserted.model_call_context_id, - inserted.served_provider_model_slug, inserted.stop_reason, - context.provider_response_id, - inserted.provider_replay, - context.input_tokens_total, context.uncached_input_tokens, - context.cache_read_input_tokens, context.cache_write_input_tokens, - context.output_tokens_total, context.reasoning_output_tokens, - inserted.created_at -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id -JOIN model_call_contexts context ON context.project_id = agent.project_id - AND context.agent_id = inserted.agent_id - AND context.id = inserted.model_call_context_id -` - -type InsertModelOutputAuthorityParams struct { - ServedProviderModelSlug string - StopReason string - ProviderReplay *json.RawMessage - ModelCallContextID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type InsertModelOutputAuthorityRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - ModelCallContextID uuid.UUID - ServedProviderModelSlug string - StopReason string - ProviderResponseID string - ProviderReplay *json.RawMessage - InputTokensTotal *int32 - UncachedInputTokens *int32 - CacheReadInputTokens *int32 - CacheWriteInputTokens *int32 - OutputTokensTotal *int32 - ReasoningOutputTokens *int32 - CreatedAt time.Time -} - -func (q *Queries) InsertModelOutputAuthority(ctx context.Context, arg InsertModelOutputAuthorityParams) (InsertModelOutputAuthorityRow, error) { - row := q.db.QueryRow(ctx, insertModelOutputAuthority, - arg.ServedProviderModelSlug, - arg.StopReason, - arg.ProviderReplay, - arg.ModelCallContextID, - arg.ProjectID, - arg.AgentID, - ) - var i InsertModelOutputAuthorityRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.ModelCallContextID, - &i.ServedProviderModelSlug, - &i.StopReason, - &i.ProviderResponseID, - &i.ProviderReplay, - &i.InputTokensTotal, - &i.UncachedInputTokens, - &i.CacheReadInputTokens, - &i.CacheWriteInputTokens, - &i.OutputTokensTotal, - &i.ReasoningOutputTokens, - &i.CreatedAt, - ) - return i, err -} - -const insertToolCallResultAuthority = `-- name: InsertToolCallResultAuthority :one -WITH inserted AS ( - INSERT INTO tool_call_results( - agent_id, tool_call_id, outcome, completed_at - ) - SELECT tool_call.agent_id, tool_call.id, - $1, statement_timestamp() - FROM tool_call_read_projection tool_call - WHERE tool_call.project_id = $2 - AND tool_call.agent_id = $3 - AND tool_call.id = $4 - AND tool_call.state = 'completed' - ON CONFLICT (agent_id, tool_call_id) DO NOTHING - RETURNING id, agent_id, tool_call_id, outcome, completed_at -) -SELECT inserted.id, agent.project_id, inserted.agent_id, - source_event.turn_id, inserted.tool_call_id, inserted.outcome, - inserted.completed_at -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id -JOIN tool_calls tool_call ON tool_call.agent_id = inserted.agent_id - AND tool_call.id = inserted.tool_call_id -JOIN agent_events source_event ON source_event.agent_id = tool_call.agent_id - AND source_event.model_output_id = tool_call.model_output_id - AND source_event.event_kind = 'model_output' -` - -type InsertToolCallResultAuthorityParams struct { - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID - ToolCallID uuid.UUID -} - -type InsertToolCallResultAuthorityRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - ToolCallID uuid.UUID - Outcome string - CompletedAt time.Time -} - -func (q *Queries) InsertToolCallResultAuthority(ctx context.Context, arg InsertToolCallResultAuthorityParams) (InsertToolCallResultAuthorityRow, error) { - row := q.db.QueryRow(ctx, insertToolCallResultAuthority, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - arg.ToolCallID, - ) - var i InsertToolCallResultAuthorityRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.ToolCallID, - &i.Outcome, - &i.CompletedAt, - ) - return i, err -} - -const insertTypedAgentEvent = `-- name: InsertTypedAgentEvent :one -WITH event_time AS ( - SELECT statement_timestamp() AS created_at - FROM agent_inputs input - WHERE input.project_id = $12 - AND input.agent_id = $2 - AND input.id = $7::uuid - - UNION ALL - - SELECT output.created_at - FROM model_outputs output - JOIN agents agent ON agent.id = output.agent_id - WHERE agent.project_id = $12 - AND output.agent_id = $2 - AND output.id = $8::uuid - - UNION ALL - - SELECT result.completed_at - FROM tool_call_results result - JOIN agents agent ON agent.id = result.agent_id - WHERE agent.project_id = $12 - AND result.agent_id = $2 - AND result.id = $9::uuid - - UNION ALL - - SELECT checkpoint.created_at - FROM context_checkpoints checkpoint - JOIN agents agent ON agent.id = checkpoint.agent_id - WHERE agent.project_id = $12 - AND checkpoint.agent_id = $2 - AND checkpoint.id = $10::uuid -) -INSERT INTO agent_events( - id, agent_id, turn_id, sequence, event_kind, - idempotency_key, agent_input_id, model_output_id, - tool_call_result_id, context_checkpoint_id, is_opening_event, created_at -) -SELECT - coalesce($1::uuid, uuidv7()), - $2, $3, $4, - $5, - $6, - $7::uuid, $8::uuid, - $9::uuid, $10::uuid, - $11, event_time.created_at -FROM event_time -ON CONFLICT (agent_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING -RETURNING id, agent_id, turn_id, is_opening_event, sequence, event_kind, created_at, coalesce(idempotency_key, '') AS idempotency_key, agent_input_id, model_output_id, tool_call_result_id, context_checkpoint_id -` - -type InsertTypedAgentEventParams struct { - ID *uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - Sequence int64 - EventKind string - IdempotencyKey *string - AgentInputID *uuid.UUID - ModelOutputID *uuid.UUID - ToolCallResultID *uuid.UUID - ContextCheckpointID *uuid.UUID - IsOpeningEvent bool - ProjectID uuid.UUID -} - -type InsertTypedAgentEventRow struct { - ID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - IsOpeningEvent bool - Sequence int64 - EventKind string - CreatedAt time.Time - IdempotencyKey string - AgentInputID *uuid.UUID - ModelOutputID *uuid.UUID - ToolCallResultID *uuid.UUID - ContextCheckpointID *uuid.UUID -} - -func (q *Queries) InsertTypedAgentEvent(ctx context.Context, arg InsertTypedAgentEventParams) (InsertTypedAgentEventRow, error) { - row := q.db.QueryRow(ctx, insertTypedAgentEvent, - arg.ID, - arg.AgentID, - arg.TurnID, - arg.Sequence, - arg.EventKind, - arg.IdempotencyKey, - arg.AgentInputID, - arg.ModelOutputID, - arg.ToolCallResultID, - arg.ContextCheckpointID, - arg.IsOpeningEvent, - arg.ProjectID, - ) - var i InsertTypedAgentEventRow - err := row.Scan( - &i.ID, - &i.AgentID, - &i.TurnID, - &i.IsOpeningEvent, - &i.Sequence, - &i.EventKind, - &i.CreatedAt, - &i.IdempotencyKey, - &i.AgentInputID, - &i.ModelOutputID, - &i.ToolCallResultID, - &i.ContextCheckpointID, - ) - return i, err -} - const listContentBlocksForAgentInputs = `-- name: ListContentBlocksForAgentInputs :many SELECT block.id, agent.project_id, block.agent_id, block.owner_kind, block.owner_agent_input_id, block.owner_model_output_id, @@ -962,253 +258,6 @@ func (q *Queries) ListContentBlocksForAgentInputs(ctx context.Context, arg ListC return items, nil } -const listContentBlocksForModelOutput = `-- name: ListContentBlocksForModelOutput :many -SELECT block.id, agent.project_id, block.agent_id, block.owner_kind, - block.owner_agent_input_id, block.owner_model_output_id, - block.owner_tool_call_result_id, block.ordinal, block.block_kind, - coalesce(block.text_content, '') AS text_content, block.structured_data, - block.artifact_id, block.tool_call_id, block.created_at -FROM content_blocks block -JOIN agents agent ON agent.id = block.agent_id -WHERE agent.project_id = $1 - AND block.agent_id = $2 - AND block.owner_model_output_id = $3 -ORDER BY block.ordinal ASC, block.id ASC -` - -type ListContentBlocksForModelOutputParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelOutputID *uuid.UUID -} - -type ListContentBlocksForModelOutputRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - OwnerKind string - OwnerAgentInputID *uuid.UUID - OwnerModelOutputID *uuid.UUID - OwnerToolCallResultID *uuid.UUID - Ordinal int32 - BlockKind string - TextContent string - StructuredData *json.RawMessage - ArtifactID *uuid.UUID - ToolCallID *uuid.UUID - CreatedAt time.Time -} - -func (q *Queries) ListContentBlocksForModelOutput(ctx context.Context, arg ListContentBlocksForModelOutputParams) ([]ListContentBlocksForModelOutputRow, error) { - rows, err := q.db.Query(ctx, listContentBlocksForModelOutput, arg.ProjectID, arg.AgentID, arg.ModelOutputID) - if err != nil { - return nil, err - } - defer rows.Close() - items := []ListContentBlocksForModelOutputRow{} - for rows.Next() { - var i ListContentBlocksForModelOutputRow - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.OwnerKind, - &i.OwnerAgentInputID, - &i.OwnerModelOutputID, - &i.OwnerToolCallResultID, - &i.Ordinal, - &i.BlockKind, - &i.TextContent, - &i.StructuredData, - &i.ArtifactID, - &i.ToolCallID, - &i.CreatedAt, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const listToolCallResultContentBlocks = `-- name: ListToolCallResultContentBlocks :many -SELECT block.id, agent.project_id, block.agent_id, block.owner_kind, - block.owner_agent_input_id, block.owner_model_output_id, - block.owner_tool_call_result_id, block.ordinal, block.block_kind, - coalesce(block.text_content, '') AS text_content, block.structured_data, - block.artifact_id, block.tool_call_id, block.exclude_from_model_context, - block.metadata, block.created_at -FROM content_blocks block -JOIN agents agent ON agent.id = block.agent_id -WHERE agent.project_id = $1 - AND block.agent_id = $2 - AND block.owner_tool_call_result_id = $3 -ORDER BY block.ordinal -` - -type ListToolCallResultContentBlocksParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ToolCallResultID *uuid.UUID -} - -type ListToolCallResultContentBlocksRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - OwnerKind string - OwnerAgentInputID *uuid.UUID - OwnerModelOutputID *uuid.UUID - OwnerToolCallResultID *uuid.UUID - Ordinal int32 - BlockKind string - TextContent string - StructuredData *json.RawMessage - ArtifactID *uuid.UUID - ToolCallID *uuid.UUID - ExcludeFromModelContext bool - Metadata json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) ListToolCallResultContentBlocks(ctx context.Context, arg ListToolCallResultContentBlocksParams) ([]ListToolCallResultContentBlocksRow, error) { - rows, err := q.db.Query(ctx, listToolCallResultContentBlocks, arg.ProjectID, arg.AgentID, arg.ToolCallResultID) - if err != nil { - return nil, err - } - defer rows.Close() - items := []ListToolCallResultContentBlocksRow{} - for rows.Next() { - var i ListToolCallResultContentBlocksRow - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.OwnerKind, - &i.OwnerAgentInputID, - &i.OwnerModelOutputID, - &i.OwnerToolCallResultID, - &i.Ordinal, - &i.BlockKind, - &i.TextContent, - &i.StructuredData, - &i.ArtifactID, - &i.ToolCallID, - &i.ExcludeFromModelContext, - &i.Metadata, - &i.CreatedAt, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const resolveConfigChangeAgentInput = `-- name: ResolveConfigChangeAgentInput :one -WITH causal_event AS MATERIALIZED ( - SELECT input.agent_config_id, event.created_at - FROM agent_inputs input - JOIN agent_events event ON event.agent_id = input.agent_id - AND event.agent_input_id = input.id - WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.input_kind = 'config_change' - AND input.state = 'received' - AND event.id = $4 - AND event.event_kind = 'agent_input' -), -resolved AS ( - UPDATE agent_inputs input - SET state = 'resolved', - admitted_event_id = $4, - admitted_at = causal_event.created_at, - resolved_at = causal_event.created_at - FROM causal_event - WHERE input.project_id = $1 - AND input.agent_id = $2 - AND input.id = $3 - AND input.state = 'received' - RETURNING input.admitted_at, input.resolved_at, input.agent_config_id -), -updated_agent AS ( - UPDATE agents agent - SET current_config_id = resolved.agent_config_id, - updated_at = resolved.resolved_at - FROM resolved - WHERE agent.project_id = $1 - AND agent.id = $2 - RETURNING agent.id -) -SELECT resolved.admitted_at, resolved.resolved_at -FROM resolved -JOIN updated_agent ON true -` - -type ResolveConfigChangeAgentInputParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - EventID uuid.UUID -} - -type ResolveConfigChangeAgentInputRow struct { - AdmittedAt *time.Time - ResolvedAt *time.Time -} - -func (q *Queries) ResolveConfigChangeAgentInput(ctx context.Context, arg ResolveConfigChangeAgentInputParams) (ResolveConfigChangeAgentInputRow, error) { - row := q.db.QueryRow(ctx, resolveConfigChangeAgentInput, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.EventID, - ) - var i ResolveConfigChangeAgentInputRow - err := row.Scan(&i.AdmittedAt, &i.ResolvedAt) - return i, err -} - -const resolveControlAgentInput = `-- name: ResolveControlAgentInput :exec -UPDATE agent_inputs -SET state = 'resolved', - admitted_event_id = $1, - admitted_at = statement_timestamp(), - resolved_at = statement_timestamp() -WHERE project_id = $2 - AND agent_id = $3 - AND id = $4 - AND input_kind = 'control' - AND control_type = $5 - AND state = 'received' -` - -type ResolveControlAgentInputParams struct { - EventID *uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - ControlType *string -} - -func (q *Queries) ResolveControlAgentInput(ctx context.Context, arg ResolveControlAgentInputParams) error { - _, err := q.db.Exec(ctx, resolveControlAgentInput, - arg.EventID, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.ControlType, - ) - return err -} - const toolCallResultHasTypedEvent = `-- name: ToolCallResultHasTypedEvent :one SELECT EXISTS ( SELECT 1 diff --git a/internal/storage/internal/dbsqlc/model_contexts.sql.go b/internal/storage/internal/dbsqlc/model_contexts.sql.go index 9db5d79af..faa501cea 100644 --- a/internal/storage/internal/dbsqlc/model_contexts.sql.go +++ b/internal/storage/internal/dbsqlc/model_contexts.sql.go @@ -13,258 +13,6 @@ import ( "github.com/google/uuid" ) -const agentHasLiveModelCallContextBeforeFrontier = `-- name: AgentHasLiveModelCallContextBeforeFrontier :one -SELECT EXISTS ( - SELECT 1 - FROM model_call_contexts context - WHERE context.project_id = $1 - AND context.agent_id = $2 - AND context.state = 'started' - AND context.input_event_sequence < $3 -)::boolean -` - -type AgentHasLiveModelCallContextBeforeFrontierParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - InputEventSequence int64 -} - -func (q *Queries) AgentHasLiveModelCallContextBeforeFrontier(ctx context.Context, arg AgentHasLiveModelCallContextBeforeFrontierParams) (bool, error) { - row := q.db.QueryRow(ctx, agentHasLiveModelCallContextBeforeFrontier, arg.ProjectID, arg.AgentID, arg.InputEventSequence) - var column_1 bool - err := row.Scan(&column_1) - return column_1, err -} - -const agentHasLiveModelCallContexts = `-- name: AgentHasLiveModelCallContexts :one -SELECT EXISTS ( - SELECT 1 - FROM model_call_contexts context - WHERE context.project_id = $1 - AND context.agent_id = $2 - AND context.state = 'started' -)::boolean -` - -type AgentHasLiveModelCallContextsParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) AgentHasLiveModelCallContexts(ctx context.Context, arg AgentHasLiveModelCallContextsParams) (bool, error) { - row := q.db.QueryRow(ctx, agentHasLiveModelCallContexts, arg.ProjectID, arg.AgentID) - var column_1 bool - err := row.Scan(&column_1) - return column_1, err -} - -const cancelRuntimeModelCallContextsForLifecycle = `-- name: CancelRuntimeModelCallContextsForLifecycle :execrows -UPDATE model_call_contexts context -SET state = 'canceled', - recovery_kind = NULL, - error_kind = 'canceled', - error_code = $1, - error_message = $2, - error_details = $3, - retry_at = NULL, - completed_at = statement_timestamp() -WHERE context.project_id = $4 - AND context.agent_id = $5 - AND context.runtime_lock_id = $6 - AND context.state = 'started' -` - -type CancelRuntimeModelCallContextsForLifecycleParams struct { - ErrorCode string - ErrorMessage string - ErrorDetails json.RawMessage - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID -} - -func (q *Queries) CancelRuntimeModelCallContextsForLifecycle(ctx context.Context, arg CancelRuntimeModelCallContextsForLifecycleParams) (int64, error) { - result, err := q.db.Exec(ctx, cancelRuntimeModelCallContextsForLifecycle, - arg.ErrorCode, - arg.ErrorMessage, - arg.ErrorDetails, - arg.ProjectID, - arg.AgentID, - arg.RuntimeLockID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const finishModelCallContext = `-- name: FinishModelCallContext :one -WITH agent_scope AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $21 - AND agent.id = $22 -), -runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agent_scope agent ON agent.id = runtime_lock.agent_id - WHERE runtime_lock.id = $23 - AND ( - $24::boolean - OR ( - runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() - ) - ) -) -UPDATE model_call_contexts context -SET state = $1, - recovery_kind = $2, - api_format = $3, - api_variant = $4, - provider_request_id = $5, - provider_response_id = $6, - error_kind = $7, - error_code = $8, - error_message = $9, - error_details = $10, - retry_at = CASE - WHEN $11::bigint IS NULL THEN NULL - ELSE statement_timestamp() + - ($11::bigint * interval '1 microsecond') - END, - input_tokens_total = $12::integer, - uncached_input_tokens = $13::integer, - cache_read_input_tokens = $14::integer, - cache_write_input_tokens = $15::integer, - output_tokens_total = $16::integer, - reasoning_output_tokens = $17::integer, - provider_reported_cost_usd = $18::text::numeric, - provider_metadata = $19, - completed_at = statement_timestamp() -FROM runtime -WHERE context.project_id = runtime.project_id - AND context.agent_id = runtime.agent_id - AND context.id = $20 - AND context.runtime_lock_id = runtime.id - AND context.state = 'started' -RETURNING context.id -` - -type FinishModelCallContextParams struct { - ToState string - RecoveryKind *string - ApiFormat string - ApiVariant string - ProviderRequestID string - ProviderResponseID string - ErrorKind string - ErrorCode string - ErrorMessage string - ErrorDetails json.RawMessage - RetryDelayMicroseconds *int64 - InputTokensTotal *int32 - UncachedInputTokens *int32 - CacheReadInputTokens *int32 - CacheWriteInputTokens *int32 - OutputTokensTotal *int32 - ReasoningOutputTokens *int32 - ProviderReportedCostUsd *string - ProviderMetadata json.RawMessage - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID - AllowInactiveRuntimeLockForTeardown bool -} - -func (q *Queries) FinishModelCallContext(ctx context.Context, arg FinishModelCallContextParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, finishModelCallContext, - arg.ToState, - arg.RecoveryKind, - arg.ApiFormat, - arg.ApiVariant, - arg.ProviderRequestID, - arg.ProviderResponseID, - arg.ErrorKind, - arg.ErrorCode, - arg.ErrorMessage, - arg.ErrorDetails, - arg.RetryDelayMicroseconds, - arg.InputTokensTotal, - arg.UncachedInputTokens, - arg.CacheReadInputTokens, - arg.CacheWriteInputTokens, - arg.OutputTokensTotal, - arg.ReasoningOutputTokens, - arg.ProviderReportedCostUsd, - arg.ProviderMetadata, - arg.ID, - arg.ProjectID, - arg.AgentID, - arg.RuntimeLockID, - arg.AllowInactiveRuntimeLockForTeardown, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const getCompactionModelCallContextByIdentity = `-- name: GetCompactionModelCallContextByIdentity :one -SELECT context.id -FROM model_call_contexts context -WHERE context.project_id = $1 - AND context.agent_id = $2 - AND context.operation_kind = 'compaction' - AND context.input_event_sequence = $3 - AND context.source_event_sequence_end = $4 -ORDER BY context.attempt_number DESC -LIMIT 1 -` - -type GetCompactionModelCallContextByIdentityParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - InputEventSequence int64 - SourceEventSequenceEnd *int64 -} - -func (q *Queries) GetCompactionModelCallContextByIdentity(ctx context.Context, arg GetCompactionModelCallContextByIdentityParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, getCompactionModelCallContextByIdentity, - arg.ProjectID, - arg.AgentID, - arg.InputEventSequence, - arg.SourceEventSequenceEnd, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const getLiveModelCallContextForRuntime = `-- name: GetLiveModelCallContextForRuntime :one -SELECT id -FROM model_call_contexts -WHERE project_id = $1 - AND agent_id = $2 - AND runtime_lock_id = $3 - AND state = 'started' -` - -type GetLiveModelCallContextForRuntimeParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID -} - -func (q *Queries) GetLiveModelCallContextForRuntime(ctx context.Context, arg GetLiveModelCallContextForRuntimeParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, getLiveModelCallContextForRuntime, arg.ProjectID, arg.AgentID, arg.RuntimeLockID) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - const getModelCallContext = `-- name: GetModelCallContext :one SELECT context.id, context.org_id, context.project_id, context.agent_id, context.operation_kind, context.attempt_number, context.agent_config_id, @@ -367,87 +115,6 @@ func (q *Queries) GetModelCallContext(ctx context.Context, arg GetModelCallConte return i, err } -const getModelCallContextProviderModelSlug = `-- name: GetModelCallContextProviderModelSlug :one -SELECT revision.provider_model_slug -FROM model_call_contexts context -JOIN configured_model_revisions revision ON revision.org_id = context.org_id - AND revision.id = context.configured_model_revision_id -WHERE context.project_id = $1 - AND context.agent_id = $2 - AND context.id = $3 -` - -type GetModelCallContextProviderModelSlugParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelCallContextID uuid.UUID -} - -func (q *Queries) GetModelCallContextProviderModelSlug(ctx context.Context, arg GetModelCallContextProviderModelSlugParams) (string, error) { - row := q.db.QueryRow(ctx, getModelCallContextProviderModelSlug, arg.ProjectID, arg.AgentID, arg.ModelCallContextID) - var provider_model_slug string - err := row.Scan(&provider_model_slug) - return provider_model_slug, err -} - -const getModelCallContextTurnID = `-- name: GetModelCallContextTurnID :one -SELECT turn_id -FROM model_call_context_turns -WHERE project_id = $1 - AND agent_id = $2 - AND model_call_context_id = $3 -` - -type GetModelCallContextTurnIDParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelCallContextID uuid.UUID -} - -func (q *Queries) GetModelCallContextTurnID(ctx context.Context, arg GetModelCallContextTurnIDParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, getModelCallContextTurnID, arg.ProjectID, arg.AgentID, arg.ModelCallContextID) - var turn_id uuid.UUID - err := row.Scan(&turn_id) - return turn_id, err -} - -const getModelCallRevisionForClaim = `-- name: GetModelCallRevisionForClaim :one -SELECT configured_model.current_revision_id, - CASE - WHEN provider.management_kind = 'cluster' - THEN COALESCE(admission.new_managed_work_allowed, true) - ELSE true - END AS new_managed_work_allowed -FROM agent_configs config -JOIN configured_models configured_model ON configured_model.org_id = config.org_id - AND configured_model.id = config.configured_model_id -JOIN model_provider_configs provider ON provider.org_id = configured_model.org_id - AND provider.id = configured_model.model_provider_config_id -LEFT JOIN org_managed_work_admission admission ON admission.org_id = config.org_id -WHERE config.project_id = $1 - AND config.id = $2 -` - -type GetModelCallRevisionForClaimParams struct { - ProjectID uuid.UUID - AgentConfigID uuid.UUID -} - -type GetModelCallRevisionForClaimRow struct { - CurrentRevisionID uuid.UUID - NewManagedWorkAllowed bool -} - -// @sqlc-vet-disable configured-models-deleted-at -// @sqlc-vet-disable model-provider-configs-deleted-at -// Model-call lineage must survive soft deletion so the unavailable model can fail durably. -func (q *Queries) GetModelCallRevisionForClaim(ctx context.Context, arg GetModelCallRevisionForClaimParams) (GetModelCallRevisionForClaimRow, error) { - row := q.db.QueryRow(ctx, getModelCallRevisionForClaim, arg.ProjectID, arg.AgentConfigID) - var i GetModelCallRevisionForClaimRow - err := row.Scan(&i.CurrentRevisionID, &i.NewManagedWorkAllowed) - return i, err -} - const getNormalModelCallContextByIdentity = `-- name: GetNormalModelCallContextByIdentity :one SELECT context.id FROM model_call_contexts context @@ -520,464 +187,3 @@ func (q *Queries) GetProviderReplaySuppressionCutoff(ctx context.Context, arg Ge err := row.Scan(&cutoff_event_sequence) return cutoff_event_sequence, err } - -const insertNextModelCallContext = `-- name: InsertNextModelCallContext :one -WITH agent_scope AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $1 - AND agent.id = $2 -), -predecessor AS MATERIALIZED ( - SELECT context.id, context.org_id, context.project_id, context.agent_id, - context.operation_kind, context.attempt_number, - context.agent_config_id, context.input_event_sequence, - context.source_event_sequence_end - FROM model_call_contexts context - JOIN agent_scope agent ON agent.project_id = context.project_id - AND agent.id = context.agent_id - WHERE context.id = $3 - AND context.state = 'failed' - AND context.recovery_kind = 'retry' - AND context.retry_at <= statement_timestamp() - AND context.attempt_number <= $4::integer - AND EXISTS ( - SELECT 1 - FROM agent_continuable_model_contexts(context.project_id, context.agent_id) continuable - WHERE continuable.model_call_context_id = context.id - AND NOT continuable.has_later_semantic_event - ) -), -selected_model AS MATERIALIZED ( - SELECT predecessor.id AS predecessor_context_id, - revision.id AS configured_model_revision_id - FROM predecessor - JOIN agent_configs config ON config.project_id = predecessor.project_id - AND config.id = predecessor.agent_config_id - JOIN configured_model_revisions revision ON revision.org_id = config.org_id - AND revision.configured_model_id = config.configured_model_id - AND revision.id = $5 -), -live_runtime AS MATERIALIZED ( - SELECT runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN predecessor context ON context.agent_id = runtime_lock.agent_id - WHERE runtime_lock.id = $6 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, attempt_number, - agent_config_id, configured_model_revision_id, input_event_sequence, - source_event_sequence_end, runtime_lock_id, state, created_at -) -SELECT predecessor.org_id, - predecessor.project_id, - predecessor.agent_id, - predecessor.operation_kind, - predecessor.attempt_number + 1, - predecessor.agent_config_id, - selected_model.configured_model_revision_id, - predecessor.input_event_sequence, - predecessor.source_event_sequence_end, - live_runtime.id, - 'started', - statement_timestamp() -FROM predecessor -JOIN selected_model ON selected_model.predecessor_context_id = predecessor.id -JOIN live_runtime ON true -RETURNING id -` - -type InsertNextModelCallContextParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - PredecessorModelCallContextID uuid.UUID - MaxRetries int32 - ConfiguredModelRevisionID uuid.UUID - RuntimeLockID uuid.UUID -} - -func (q *Queries) InsertNextModelCallContext(ctx context.Context, arg InsertNextModelCallContextParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, insertNextModelCallContext, - arg.ProjectID, - arg.AgentID, - arg.PredecessorModelCallContextID, - arg.MaxRetries, - arg.ConfiguredModelRevisionID, - arg.RuntimeLockID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const insertNormalModelCallContext = `-- name: InsertNormalModelCallContext :one -WITH active_agent AS MATERIALIZED ( - SELECT agent.org_id, agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $2 - AND agent.id = $3 - AND agent.state = 'active' -), -live_runtime AS MATERIALIZED ( - SELECT runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN active_agent agent ON agent.id = runtime_lock.agent_id - WHERE runtime_lock.id = $4 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -), -active_config AS MATERIALIZED ( - SELECT input.agent_config_id - FROM agent_events event - JOIN agent_inputs input ON input.agent_id = event.agent_id - AND input.id = event.agent_input_id - JOIN active_agent agent ON agent.id = event.agent_id - WHERE event.sequence <= $1 - AND event.event_kind = 'agent_input' - AND input.input_kind = 'config_change' - AND input.state = 'resolved' - AND input.admitted_event_id = event.id - AND input.agent_config_id IS NOT NULL - ORDER BY event.sequence DESC - LIMIT 1 -), -selected_model AS MATERIALIZED ( - SELECT config.id AS agent_config_id, - revision.id AS configured_model_revision_id - FROM agent_configs config - JOIN active_config ON active_config.agent_config_id = config.id - JOIN configured_model_revisions revision ON revision.org_id = config.org_id - AND revision.configured_model_id = config.configured_model_id - AND revision.id = $5 - WHERE config.project_id = $2 -) -INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, attempt_number, - agent_config_id, configured_model_revision_id, input_event_sequence, - runtime_lock_id, state, created_at -) -SELECT active_agent.org_id, - active_agent.project_id, - active_agent.id, - 'normal', - 1, - selected_model.agent_config_id, - selected_model.configured_model_revision_id, - $1, - live_runtime.id, - 'started', - statement_timestamp() -FROM active_agent -JOIN live_runtime ON true -JOIN selected_model ON true -RETURNING id -` - -type InsertNormalModelCallContextParams struct { - InputEventSequence int64 - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID - ConfiguredModelRevisionID uuid.UUID -} - -func (q *Queries) InsertNormalModelCallContext(ctx context.Context, arg InsertNormalModelCallContextParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, insertNormalModelCallContext, - arg.InputEventSequence, - arg.ProjectID, - arg.AgentID, - arg.RuntimeLockID, - arg.ConfiguredModelRevisionID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const insertTriggeredCompactionModelCallContext = `-- name: InsertTriggeredCompactionModelCallContext :one -WITH active_agent AS MATERIALIZED ( - SELECT agent.org_id, agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $2 - AND agent.id = $3 - AND agent.state = 'active' -), -live_runtime AS MATERIALIZED ( - SELECT runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN active_agent agent ON agent.id = runtime_lock.agent_id - WHERE runtime_lock.id = $4 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -), -latest_compaction AS MATERIALIZED ( - SELECT dependency.id, dependency.state, dependency.recovery_kind, - dependency.input_event_sequence, dependency.source_event_sequence_end - FROM model_call_contexts dependency - JOIN model_call_contexts parent ON parent.project_id = dependency.project_id - AND parent.agent_id = dependency.agent_id - AND parent.id = $5 - JOIN active_agent agent ON agent.project_id = parent.project_id - AND agent.id = parent.agent_id - WHERE dependency.operation_kind = 'compaction' - AND dependency.input_event_sequence = parent.input_event_sequence - ORDER BY dependency.source_event_sequence_end ASC, - dependency.attempt_number DESC - LIMIT 1 -), -blocked AS MATERIALIZED ( - SELECT context.id, context.agent_config_id, - context.input_event_sequence - FROM model_call_contexts context - JOIN active_agent agent ON agent.project_id = context.project_id - AND agent.id = context.agent_id - LEFT JOIN latest_compaction latest ON true - WHERE context.id = $5 - AND context.operation_kind = 'normal' - AND context.state = 'failed' - AND context.recovery_kind = 'compact' - AND ( - latest.id IS NULL - OR ( - latest.state = 'failed' - AND latest.recovery_kind = 'reduce_compaction_source' - AND latest.input_event_sequence = context.input_event_sequence - AND $1::bigint < latest.source_event_sequence_end - ) - ) -), -selected_model AS MATERIALIZED ( - SELECT blocked.id AS blocked_context_id, - revision.id AS configured_model_revision_id - FROM blocked - JOIN agent_configs config ON config.project_id = $2 - AND config.id = blocked.agent_config_id - JOIN configured_model_revisions revision ON revision.org_id = config.org_id - AND revision.configured_model_id = config.configured_model_id - AND revision.id = $6 -) -INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, attempt_number, - agent_config_id, configured_model_revision_id, input_event_sequence, - source_event_sequence_end, runtime_lock_id, state, created_at -) -SELECT active_agent.org_id, - active_agent.project_id, - active_agent.id, - 'compaction', - 1, - blocked.agent_config_id, - selected_model.configured_model_revision_id, - blocked.input_event_sequence, - $1, - live_runtime.id, - 'started', - statement_timestamp() -FROM active_agent -JOIN live_runtime ON true -JOIN blocked ON true -JOIN selected_model ON selected_model.blocked_context_id = blocked.id -RETURNING id -` - -type InsertTriggeredCompactionModelCallContextParams struct { - SourceEventSequenceEnd *int64 - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID - ParentModelCallContextID uuid.UUID - ConfiguredModelRevisionID uuid.UUID -} - -func (q *Queries) InsertTriggeredCompactionModelCallContext(ctx context.Context, arg InsertTriggeredCompactionModelCallContextParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, insertTriggeredCompactionModelCallContext, - arg.SourceEventSequenceEnd, - arg.ProjectID, - arg.AgentID, - arg.RuntimeLockID, - arg.ParentModelCallContextID, - arg.ConfiguredModelRevisionID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const interruptRuntimeModelCallContextForRetry = `-- name: InterruptRuntimeModelCallContextForRetry :execrows -UPDATE model_call_contexts context -SET state = 'failed', - recovery_kind = 'retry', - error_kind = $1, - error_code = $2, - error_message = $3, - error_details = $4::jsonb, - retry_at = statement_timestamp() + - ($5::bigint * interval '1 microsecond'), - completed_at = statement_timestamp() -WHERE context.project_id = $6 - AND context.agent_id = $7 - AND context.id = $8 - AND context.runtime_lock_id = $9 - AND context.state = 'started' -` - -type InterruptRuntimeModelCallContextForRetryParams struct { - ErrorKind string - ErrorCode string - ErrorMessage string - ErrorDetails json.RawMessage - RetryDelayMicroseconds int64 - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID - RuntimeLockID uuid.UUID -} - -func (q *Queries) InterruptRuntimeModelCallContextForRetry(ctx context.Context, arg InterruptRuntimeModelCallContextForRetryParams) (int64, error) { - result, err := q.db.Exec(ctx, interruptRuntimeModelCallContextForRetry, - arg.ErrorKind, - arg.ErrorCode, - arg.ErrorMessage, - arg.ErrorDetails, - arg.RetryDelayMicroseconds, - arg.ProjectID, - arg.AgentID, - arg.ID, - arg.RuntimeLockID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const listModelCallOpeningContentInputs = `-- name: ListModelCallOpeningContentInputs :many -SELECT opening.input_id::uuid AS input_id, - opening.event_sequence::bigint AS event_sequence -FROM agent_model_call_opening_content_inputs( - $1, - $2, - $3, - $4 -) AS opening(input_id, event_sequence) -ORDER BY opening.event_sequence -` - -type ListModelCallOpeningContentInputsParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - InputEventSequence int64 -} - -type ListModelCallOpeningContentInputsRow struct { - InputID uuid.UUID - EventSequence int64 -} - -func (q *Queries) ListModelCallOpeningContentInputs(ctx context.Context, arg ListModelCallOpeningContentInputsParams) ([]ListModelCallOpeningContentInputsRow, error) { - rows, err := q.db.Query(ctx, listModelCallOpeningContentInputs, - arg.ProjectID, - arg.AgentID, - arg.TurnID, - arg.InputEventSequence, - ) - if err != nil { - return nil, err - } - defer rows.Close() - items := []ListModelCallOpeningContentInputsRow{} - for rows.Next() { - var i ListModelCallOpeningContentInputsRow - if err := rows.Scan(&i.InputID, &i.EventSequence); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const modelCallContextHasLaterSemanticEvent = `-- name: ModelCallContextHasLaterSemanticEvent :one -SELECT model_call_context_has_later_semantic_event( - $1, - $2, - $3 -) AS has_later_semantic_event -` - -type ModelCallContextHasLaterSemanticEventParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelCallContextID uuid.UUID -} - -func (q *Queries) ModelCallContextHasLaterSemanticEvent(ctx context.Context, arg ModelCallContextHasLaterSemanticEventParams) (bool, error) { - row := q.db.QueryRow(ctx, modelCallContextHasLaterSemanticEvent, arg.ProjectID, arg.AgentID, arg.ModelCallContextID) - var has_later_semantic_event bool - err := row.Scan(&has_later_semantic_event) - return has_later_semantic_event, err -} - -const openingContentInputSetMatchesInputSequence = `-- name: OpeningContentInputSetMatchesInputSequence :one -WITH requested AS ( - SELECT input_id - FROM unnest($1::uuid[]) AS input_ids(input_id) -), -requested_distinct AS ( - SELECT DISTINCT input_id - FROM requested -), -target_turn AS MATERIALIZED ( - SELECT agent_turn_id_at_event_sequence( - $2, - $3 - ) AS turn_id -), -opening AS ( - SELECT opening.input_id - FROM target_turn turn - CROSS JOIN LATERAL agent_model_call_opening_content_inputs( - $4, - $2, - turn.turn_id, - $3 - ) AS opening(input_id, event_sequence) -) -SELECT ( - (SELECT count(*) FROM requested) = (SELECT count(*) FROM requested_distinct) - AND NOT EXISTS ( - SELECT input_id FROM requested_distinct - EXCEPT - SELECT input_id FROM opening - ) - AND NOT EXISTS ( - SELECT input_id FROM opening - EXCEPT - SELECT input_id FROM requested_distinct - ) -)::boolean -` - -type OpeningContentInputSetMatchesInputSequenceParams struct { - InputIds []uuid.UUID - AgentID uuid.UUID - InputEventSequence int64 - ProjectID uuid.UUID -} - -func (q *Queries) OpeningContentInputSetMatchesInputSequence(ctx context.Context, arg OpeningContentInputSetMatchesInputSequenceParams) (bool, error) { - row := q.db.QueryRow(ctx, openingContentInputSetMatchesInputSequence, - arg.InputIds, - arg.AgentID, - arg.InputEventSequence, - arg.ProjectID, - ) - var column_1 bool - err := row.Scan(&column_1) - return column_1, err -} diff --git a/internal/storage/internal/dbsqlc/processes.sql.go b/internal/storage/internal/dbsqlc/processes.sql.go index 783be7022..8dbbafd50 100644 --- a/internal/storage/internal/dbsqlc/processes.sql.go +++ b/internal/storage/internal/dbsqlc/processes.sql.go @@ -1151,6 +1151,72 @@ func (q *Queries) ListActiveProcesses(ctx context.Context, arg ListActiveProcess return items, nil } +const listAgentsForExecutionRevoked = `-- name: ListAgentsForExecutionRevoked :many +SELECT agent.id +FROM agents agent +WHERE agent.project_id = $1 + AND ( + ($2::uuid IS NOT NULL AND agent.id = $2::uuid) + OR ( + $2::uuid IS NULL + AND ($3::uuid IS NOT NULL OR $4::uuid IS NOT NULL) + AND EXISTS ( + SELECT 1 + FROM agent_machine_bindings binding + JOIN project_machine_grants pmgrant ON pmgrant.project_id = binding.project_id + AND pmgrant.machine_id = binding.machine_id + WHERE binding.project_id = agent.project_id + AND binding.agent_id = agent.id + AND binding.state = 'attached' + AND ( + $3::uuid IS NULL + OR pmgrant.id = $3::uuid + ) + AND ( + $4::uuid IS NULL + OR ( + pmgrant.source_kind = 'pool' + AND pmgrant.project_machine_pool_grant_id = $4::uuid + ) + ) + ) + ) + ) +ORDER BY agent.id +` + +type ListAgentsForExecutionRevokedParams struct { + ProjectID uuid.UUID + AgentID *uuid.UUID + ProjectMachineGrantID *uuid.UUID + ProjectMachinePoolGrantID *uuid.UUID +} + +func (q *Queries) ListAgentsForExecutionRevoked(ctx context.Context, arg ListAgentsForExecutionRevokedParams) ([]uuid.UUID, error) { + rows, err := q.db.Query(ctx, listAgentsForExecutionRevoked, + arg.ProjectID, + arg.AgentID, + arg.ProjectMachineGrantID, + arg.ProjectMachinePoolGrantID, + ) + if err != nil { + return nil, err + } + defer rows.Close() + items := []uuid.UUID{} + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const listDaemonProcessOffers = `-- name: ListDaemonProcessOffers :many WITH runtime AS MATERIALIZED ( SELECT runtime.org_id, runtime.machine_id @@ -1597,73 +1663,6 @@ func (q *Queries) ListProcessesForMachineReconciliation(ctx context.Context, arg return items, nil } -const lockAgentsForExecutionRevoked = `-- name: LockAgentsForExecutionRevoked :many -SELECT agent.id -FROM agents agent -WHERE agent.project_id = $1 - AND ( - ($2::uuid IS NOT NULL AND agent.id = $2::uuid) - OR ( - $2::uuid IS NULL - AND ($3::uuid IS NOT NULL OR $4::uuid IS NOT NULL) - AND EXISTS ( - SELECT 1 - FROM agent_machine_bindings binding - JOIN project_machine_grants pmgrant ON pmgrant.project_id = binding.project_id - AND pmgrant.machine_id = binding.machine_id - WHERE binding.project_id = agent.project_id - AND binding.agent_id = agent.id - AND binding.state = 'attached' - AND ( - $3::uuid IS NULL - OR pmgrant.id = $3::uuid - ) - AND ( - $4::uuid IS NULL - OR ( - pmgrant.source_kind = 'pool' - AND pmgrant.project_machine_pool_grant_id = $4::uuid - ) - ) - ) - ) - ) -ORDER BY agent.id -FOR UPDATE -` - -type LockAgentsForExecutionRevokedParams struct { - ProjectID uuid.UUID - AgentID *uuid.UUID - ProjectMachineGrantID *uuid.UUID - ProjectMachinePoolGrantID *uuid.UUID -} - -func (q *Queries) LockAgentsForExecutionRevoked(ctx context.Context, arg LockAgentsForExecutionRevokedParams) ([]uuid.UUID, error) { - rows, err := q.db.Query(ctx, lockAgentsForExecutionRevoked, - arg.ProjectID, - arg.AgentID, - arg.ProjectMachineGrantID, - arg.ProjectMachinePoolGrantID, - ) - if err != nil { - return nil, err - } - defer rows.Close() - items := []uuid.UUID{} - for rows.Next() { - var id uuid.UUID - if err := rows.Scan(&id); err != nil { - return nil, err - } - items = append(items, id) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const lockDaemonProcessForAccept = `-- name: LockDaemonProcessForAccept :one SELECT process.id FROM processes process diff --git a/internal/storage/internal/dbsqlc/secrets.sql.go b/internal/storage/internal/dbsqlc/secrets.sql.go index 147a83780..f26e32c31 100644 --- a/internal/storage/internal/dbsqlc/secrets.sql.go +++ b/internal/storage/internal/dbsqlc/secrets.sql.go @@ -189,6 +189,69 @@ func (q *Queries) DeleteSecretVersions(ctx context.Context, arg DeleteSecretVers return err } +const getCurrentSecretPayload = `-- name: GetCurrentSecretPayload :one +SELECT v.id, v.org_id, v.secret_id, v.version_number, v.payload_keys, v.encryption_scheme, + v.key_id, v.dek_wrapped_by, v.encrypted_dek, v.encrypted_dek_nonce, v.nonce, v.ciphertext, v.created_at, + (v.oauth_access_token_expires_at IS NOT NULL)::boolean AS oauth_access_token_expires, + coalesce(greatest(extract(epoch FROM v.oauth_access_token_expires_at - statement_timestamp()), 0), 0)::double precision AS oauth_access_token_remaining_seconds, + s.owner_kind, s.management_kind, s.kind +FROM secrets s +JOIN secret_versions v ON v.org_id = s.org_id AND v.secret_id = s.id AND v.id = s.current_version_id +WHERE s.org_id = $1 AND s.id = $2 AND s.deleted_at IS NULL +` + +type GetCurrentSecretPayloadParams struct { + OrgID uuid.UUID + ID uuid.UUID +} + +type GetCurrentSecretPayloadRow struct { + ID uuid.UUID + OrgID uuid.UUID + SecretID uuid.UUID + VersionNumber int32 + PayloadKeys []string + EncryptionScheme string + KeyID string + DekWrappedBy string + EncryptedDek []byte + EncryptedDekNonce []byte + Nonce []byte + Ciphertext []byte + CreatedAt time.Time + OauthAccessTokenExpires bool + OauthAccessTokenRemainingSeconds float64 + OwnerKind string + ManagementKind string + Kind string +} + +func (q *Queries) GetCurrentSecretPayload(ctx context.Context, arg GetCurrentSecretPayloadParams) (GetCurrentSecretPayloadRow, error) { + row := q.db.QueryRow(ctx, getCurrentSecretPayload, arg.OrgID, arg.ID) + var i GetCurrentSecretPayloadRow + err := row.Scan( + &i.ID, + &i.OrgID, + &i.SecretID, + &i.VersionNumber, + &i.PayloadKeys, + &i.EncryptionScheme, + &i.KeyID, + &i.DekWrappedBy, + &i.EncryptedDek, + &i.EncryptedDekNonce, + &i.Nonce, + &i.Ciphertext, + &i.CreatedAt, + &i.OauthAccessTokenExpires, + &i.OauthAccessTokenRemainingSeconds, + &i.OwnerKind, + &i.ManagementKind, + &i.Kind, + ) + return i, err +} + const getProjectAvailableSecret = `-- name: GetProjectAvailableSecret :one WITH available AS ( SELECT s.id, s.org_id, s.management_kind, s.owner_kind, s.owner_project_id, s.owner_user_id, s.name, s.kind, s.metadata, s.current_version_id, v.version_number AS current_version_number, v.payload_keys, NULL::uuid AS grant_id, s.created_at, s.updated_at diff --git a/internal/storage/internal/dbsqlc/subagents.sql.go b/internal/storage/internal/dbsqlc/subagents.sql.go index 5750a567b..5630a5fef 100644 --- a/internal/storage/internal/dbsqlc/subagents.sql.go +++ b/internal/storage/internal/dbsqlc/subagents.sql.go @@ -55,7 +55,7 @@ SELECT EXISTS ( FROM agent_wakeups wake WHERE wake.agent_id = subtree.id ) - OR agent_has_incomplete_tool_batch(subtree.project_id, subtree.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=subtree.id AND h.incomplete_tools) ) )::boolean AS idle ` @@ -218,8 +218,8 @@ SELECT agent.id, FROM agent_wakeups wake WHERE wake.agent_id = agent.id ) - OR agent_next_wakeup_ready_at(agent.project_id, agent.id) IS NOT NULL - OR agent_has_incomplete_tool_batch(agent.project_id, agent.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.logical_ready_at IS NOT NULL) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.incomplete_tools) )::boolean AS is_running FROM agents agent WHERE agent.project_id = ANY($1::uuid[]) @@ -444,8 +444,8 @@ SELECT agent.id, FROM agent_wakeups wake WHERE wake.agent_id = agent.id ) - OR agent_next_wakeup_ready_at(agent.project_id, agent.id) IS NOT NULL - OR agent_has_incomplete_tool_batch(agent.project_id, agent.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.logical_ready_at IS NOT NULL) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.incomplete_tools) )::boolean AS is_running FROM agents agent WHERE agent.project_id = $1 @@ -542,7 +542,7 @@ WHERE NOT EXISTS ( FROM agent_wakeups wake WHERE wake.agent_id = subtree.id ) - OR agent_has_incomplete_tool_batch(subtree.project_id, subtree.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=subtree.id AND h.incomplete_tools) ) ) ORDER BY candidate.created_at, candidate.id diff --git a/internal/storage/internal/dbsqlc/tool_calls.sql.go b/internal/storage/internal/dbsqlc/tool_calls.sql.go index fd40fd587..c86b0ae54 100644 --- a/internal/storage/internal/dbsqlc/tool_calls.sql.go +++ b/internal/storage/internal/dbsqlc/tool_calls.sql.go @@ -13,908 +13,16 @@ import ( "github.com/google/uuid" ) -const cancelNonTerminalToolCallsForAgent = `-- name: CancelNonTerminalToolCallsForAgent :many -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $2 - AND agent.id = $3 - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.state <> 'completed' - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id - AND projection.turn_id = $1 -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - 'canceled'::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CancelNonTerminalToolCallsForAgentParams struct { - TurnID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CancelNonTerminalToolCallsForAgentRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CancelNonTerminalToolCallsForAgent(ctx context.Context, arg CancelNonTerminalToolCallsForAgentParams) ([]CancelNonTerminalToolCallsForAgentRow, error) { - rows, err := q.db.Query(ctx, cancelNonTerminalToolCallsForAgent, arg.TurnID, arg.ProjectID, arg.AgentID) - if err != nil { - return nil, err - } - defer rows.Close() - items := []CancelNonTerminalToolCallsForAgentRow{} - for rows.Next() { - var i CancelNonTerminalToolCallsForAgentRow - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const completeCustomToolCall = `-- name: CompleteCustomToolCall :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - JOIN tool_calls tool_call ON tool_call.agent_id = agent.id - WHERE agent.project_id = $3 - AND tool_call.agent_id = $4 - AND tool_call.id = $1 - FOR UPDATE OF agent -) -UPDATE tool_calls call -SET state = 'completed' -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = $1 - AND call.state = 'ready' - AND call.runtime_lock_id IS NULL - AND call.type = 'custom' - AND $2::text IN ('succeeded', 'failed') - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $2::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteCustomToolCallParams struct { - ID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CompleteCustomToolCallRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteCustomToolCall(ctx context.Context, arg CompleteCustomToolCallParams) (CompleteCustomToolCallRow, error) { - row := q.db.QueryRow(ctx, completeCustomToolCall, - arg.ID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - ) - var i CompleteCustomToolCallRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeRuntimeToolCall = `-- name: CompleteRuntimeToolCall :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $3 - AND runtime_lock.agent_id = $4 - AND runtime_lock.id = $5 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM live_runtime runtime_lock -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = $1 - AND call.state = 'running' - AND call.runtime_lock_id = runtime_lock.id - AND call.type IN ('built_in', 'mcp') - AND projection.project_id = runtime_lock.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $2::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteRuntimeToolCallParams struct { - ID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID -} - -type CompleteRuntimeToolCallRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteRuntimeToolCall(ctx context.Context, arg CompleteRuntimeToolCallParams) (CompleteRuntimeToolCallRow, error) { - row := q.db.QueryRow(ctx, completeRuntimeToolCall, - arg.ID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - arg.RuntimeLockID, - ) - var i CompleteRuntimeToolCallRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeToolCall = `-- name: CompleteToolCall :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $5 - AND runtime_lock.agent_id = $1 - AND runtime_lock.id = $4 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM live_runtime runtime_lock -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = $1 - AND call.id = $2 - AND call.state <> 'completed' - AND ( - (call.type IN ('built_in', 'mcp') AND call.state = 'ready') - OR ( - call.state IN ('awaiting_authorization', 'awaiting_permission') - AND $3::text IN ('failed', 'denied', 'canceled') - ) - ) - AND runtime_lock.agent_id = call.agent_id - AND runtime_lock.id = $4 - AND projection.project_id = runtime_lock.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $3::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteToolCallParams struct { - AgentID uuid.UUID - ID uuid.UUID - Outcome string - RuntimeLockID uuid.UUID - ProjectID uuid.UUID -} - -type CompleteToolCallRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteToolCall(ctx context.Context, arg CompleteToolCallParams) (CompleteToolCallRow, error) { - row := q.db.QueryRow(ctx, completeToolCall, - arg.AgentID, - arg.ID, - arg.Outcome, - arg.RuntimeLockID, - arg.ProjectID, - ) - var i CompleteToolCallRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeToolCallFromPermissionInteraction = `-- name: CompleteToolCallFromPermissionInteraction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $2 - AND agent.id = $3 - FOR UPDATE -), -permission_interaction AS MATERIALIZED ( - SELECT agent.project_id, interaction.agent_id, interaction.tool_call_id - FROM agent_interactions interaction - JOIN locked_agent agent ON agent.id = interaction.agent_id - WHERE interaction.agent_id = $3 - AND interaction.id = $4 - AND interaction.interaction_kind = 'permission' - AND interaction.state IN ('resolved', 'canceled') -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM permission_interaction interaction -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = interaction.agent_id - AND call.id = interaction.tool_call_id - AND call.state = 'awaiting_permission' - AND projection.project_id = interaction.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $1::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteToolCallFromPermissionInteractionParams struct { - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID - InteractionID uuid.UUID -} - -type CompleteToolCallFromPermissionInteractionRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteToolCallFromPermissionInteraction(ctx context.Context, arg CompleteToolCallFromPermissionInteractionParams) (CompleteToolCallFromPermissionInteractionRow, error) { - row := q.db.QueryRow(ctx, completeToolCallFromPermissionInteraction, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - arg.InteractionID, - ) - var i CompleteToolCallFromPermissionInteractionRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeToolCallFromProcess = `-- name: CompleteToolCallFromProcess :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $4 - AND agent.id = $5 - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM processes tp -CROSS JOIN locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = $1 - AND call.state = 'waiting' - AND call.type = 'built_in' - AND tp.project_id = agent.project_id - AND tp.agent_id = call.agent_id - AND tp.id = $2 - AND tp.tool_call_id = call.id - AND tp.state IN ('exited', 'failed', 'killed', 'unknown') - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $3::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteToolCallFromProcessParams struct { - ID uuid.UUID - ProcessID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CompleteToolCallFromProcessRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteToolCallFromProcess(ctx context.Context, arg CompleteToolCallFromProcessParams) (CompleteToolCallFromProcessRow, error) { - row := q.db.QueryRow(ctx, completeToolCallFromProcess, - arg.ID, - arg.ProcessID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - ) - var i CompleteToolCallFromProcessRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeToolCallFromProcessAction = `-- name: CompleteToolCallFromProcessAction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $5 - AND agent.id = $6 - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM process_actions action -CROSS JOIN locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = action.tool_call_id - AND call.id = $1 - AND call.state = 'waiting' - AND call.type = 'built_in' - AND action.project_id = agent.project_id - AND action.agent_id = call.agent_id - AND action.process_id = $2 - AND action.id = $3 - AND action.state IN ('applied', 'failed', 'unknown') - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $4::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteToolCallFromProcessActionParams struct { - ToolCallID uuid.UUID - ProcessID uuid.UUID - ProcessActionID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CompleteToolCallFromProcessActionRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteToolCallFromProcessAction(ctx context.Context, arg CompleteToolCallFromProcessActionParams) (CompleteToolCallFromProcessActionRow, error) { - row := q.db.QueryRow(ctx, completeToolCallFromProcessAction, - arg.ToolCallID, - arg.ProcessID, - arg.ProcessActionID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - ) - var i CompleteToolCallFromProcessActionRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeToolCallFromQuestionInteraction = `-- name: CompleteToolCallFromQuestionInteraction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $3 - AND agent.id = $4 - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.state IN ('running', 'waiting') - AND call.type = 'built_in' - AND EXISTS ( - SELECT 1 - FROM agent_interactions question - WHERE question.agent_id = call.agent_id - AND question.tool_call_id = call.id - AND question.id = $1 - AND question.interaction_kind = 'question' - AND question.state IN ('resolved', 'canceled') - ) - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $2::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteToolCallFromQuestionInteractionParams struct { - InteractionID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CompleteToolCallFromQuestionInteractionRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteToolCallFromQuestionInteraction(ctx context.Context, arg CompleteToolCallFromQuestionInteractionParams) (CompleteToolCallFromQuestionInteractionRow, error) { - row := q.db.QueryRow(ctx, completeToolCallFromQuestionInteraction, - arg.InteractionID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - ) - var i CompleteToolCallFromQuestionInteractionRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeToolCallFromStartedProcess = `-- name: CompleteToolCallFromStartedProcess :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $4 - AND agent.id = $5 - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM processes process -CROSS JOIN locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = $1 - AND call.state = 'waiting' - AND call.type = 'built_in' - AND process.project_id = agent.project_id - AND process.agent_id = call.agent_id - AND process.id = $2 - AND process.tool_call_id = call.id - AND process.state = 'running' - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $3::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteToolCallFromStartedProcessParams struct { - ID uuid.UUID - ProcessID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CompleteToolCallFromStartedProcessRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteToolCallFromStartedProcess(ctx context.Context, arg CompleteToolCallFromStartedProcessParams) (CompleteToolCallFromStartedProcessRow, error) { - row := q.db.QueryRow(ctx, completeToolCallFromStartedProcess, - arg.ID, - arg.ProcessID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - ) - var i CompleteToolCallFromStartedProcessRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - -const completeWaitingBuiltInToolCall = `-- name: CompleteWaitingBuiltInToolCall :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $3 - AND agent.id = $4 - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = $1 - AND call.state = 'waiting' - AND call.type = 'built_in' - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - $2::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at -` - -type CompleteWaitingBuiltInToolCallParams struct { - ID uuid.UUID - Outcome string - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type CompleteWaitingBuiltInToolCallRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time -} - -func (q *Queries) CompleteWaitingBuiltInToolCall(ctx context.Context, arg CompleteWaitingBuiltInToolCallParams) (CompleteWaitingBuiltInToolCallRow, error) { - row := q.db.QueryRow(ctx, completeWaitingBuiltInToolCall, - arg.ID, - arg.Outcome, - arg.ProjectID, - arg.AgentID, - ) - var i CompleteWaitingBuiltInToolCallRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - ) - return i, err -} - const getToolCall = `-- name: GetToolCall :one +WITH selected_call AS MATERIALIZED ( + SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, + tc.model_call_context_id, tc.model_output_id, tc.provider_call_id, tc.name, tc.input, + tc.type, tc.state, tc.runtime_lock_id, tc.created_at + FROM tool_call_read_projection tc + WHERE tc.project_id = $1 + AND tc.agent_id = $2 + AND tc.id = $3 +) SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.model_call_context_id, tc.provider_call_id, tc.name, tc.input, tc.type, @@ -923,7 +31,7 @@ SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, tc.created_at, result.completed_at -FROM tool_call_read_projection tc +FROM selected_call tc LEFT JOIN tool_call_results result ON result.agent_id = tc.agent_id AND result.tool_call_id = tc.id LEFT JOIN LATERAL ( @@ -944,7 +52,6 @@ LEFT JOIN LATERAL ( WHERE block.agent_id = result.agent_id AND block.owner_tool_call_result_id = result.id ) result_blocks ON true -WHERE tc.project_id = $1 AND tc.agent_id = $2 AND tc.id = $3 ` type GetToolCallParams struct { @@ -997,6 +104,16 @@ func (q *Queries) GetToolCall(ctx context.Context, arg GetToolCallParams) (GetTo } const getToolCallByProviderCall = `-- name: GetToolCallByProviderCall :one +WITH selected_call AS MATERIALIZED ( + SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, + tc.model_call_context_id, tc.model_output_id, tc.provider_call_id, tc.name, tc.input, + tc.type, tc.state, tc.runtime_lock_id, tc.created_at + FROM tool_call_read_projection tc + WHERE tc.project_id = $1 + AND tc.agent_id = $2 + AND tc.model_call_context_id = $3 + AND tc.provider_call_id = $4 +) SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.model_call_context_id, tc.provider_call_id, tc.name, tc.input, tc.type, @@ -1005,199 +122,37 @@ SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, tc.created_at, result.completed_at -FROM tool_call_read_projection tc +FROM selected_call tc LEFT JOIN tool_call_results result ON result.agent_id = tc.agent_id AND result.tool_call_id = tc.id LEFT JOIN LATERAL ( - SELECT coalesce(jsonb_agg( - CASE - WHEN block.block_kind = 'text' THEN jsonb_build_object('type', 'text', 'text', block.text_content) - WHEN block.block_kind = 'structured_data' THEN jsonb_build_object('type', 'structured_data', 'value', block.structured_data) - WHEN block.block_kind = 'artifact' THEN - jsonb_build_object('type', 'media_ref', 'artifact_id', block.artifact_id::text) || - CASE - WHEN block.exclude_from_model_context THEN jsonb_build_object('exclude_from_model_context', true) - ELSE '{}'::jsonb - END - END - ORDER BY block.ordinal, block.id - ) FILTER (WHERE block.id IS NOT NULL AND block.block_kind IN ('text', 'structured_data', 'artifact')), '[]'::jsonb) AS content_parts - FROM content_blocks block - WHERE block.agent_id = result.agent_id - AND block.owner_tool_call_result_id = result.id -) result_blocks ON true -WHERE tc.project_id = $1 AND tc.agent_id = $2 AND tc.model_call_context_id = $3 AND tc.provider_call_id = $4 -` - -type GetToolCallByProviderCallParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string -} - -type GetToolCallByProviderCallRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time - CompletedAt *time.Time -} - -func (q *Queries) GetToolCallByProviderCall(ctx context.Context, arg GetToolCallByProviderCallParams) (GetToolCallByProviderCallRow, error) { - row := q.db.QueryRow(ctx, getToolCallByProviderCall, - arg.ProjectID, - arg.AgentID, - arg.ModelCallContextID, - arg.ProviderCallID, - ) - var i GetToolCallByProviderCallRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - &i.CompletedAt, - ) - return i, err -} - -const getToolCallDispatchState = `-- name: GetToolCallDispatchState :one -SELECT state, runtime_lock_id -FROM tool_call_read_projection -WHERE project_id = $1 - AND agent_id = $2 - AND id = $3 -` - -type GetToolCallDispatchStateParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ID uuid.UUID -} - -type GetToolCallDispatchStateRow struct { - State string - RuntimeLockID *uuid.UUID -} - -func (q *Queries) GetToolCallDispatchState(ctx context.Context, arg GetToolCallDispatchStateParams) (GetToolCallDispatchStateRow, error) { - row := q.db.QueryRow(ctx, getToolCallDispatchState, arg.ProjectID, arg.AgentID, arg.ID) - var i GetToolCallDispatchStateRow - err := row.Scan(&i.State, &i.RuntimeLockID) - return i, err -} - -const insertToolCall = `-- name: InsertToolCall :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $1 - AND runtime_lock.agent_id = $2 - AND runtime_lock.id = $3 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -), -source_event AS MATERIALIZED ( - SELECT event.id, event.turn_id, event.model_output_id, model_output.created_at - FROM agent_events event - JOIN agents agent ON agent.id = event.agent_id - JOIN model_outputs model_output ON model_output.agent_id = event.agent_id - AND model_output.id = event.model_output_id - WHERE agent.project_id = $1 - AND event.agent_id = $2 - AND event.id = $4 - AND event.event_kind = 'model_output' - AND event.model_output_id IS NOT NULL - AND model_output.model_call_context_id = $5 -), -candidate AS MATERIALIZED ( - SELECT coalesce($6::uuid, uuidv7()) AS id, - mcc.project_id, - mcc.agent_id, - mcc.id AS model_call_context_id, - source_event.id AS source_event_id, - source_event.turn_id, - source_event.model_output_id, - $7::text AS provider_call_id, - $8::text AS name, - $9::jsonb AS input, - $10::text AS type, - source_event.created_at - FROM model_call_contexts mcc - JOIN live_runtime runtime_lock ON runtime_lock.project_id = mcc.project_id - AND runtime_lock.agent_id = mcc.agent_id - AND runtime_lock.id = mcc.runtime_lock_id - JOIN source_event ON true - WHERE mcc.project_id = $1 - AND mcc.agent_id = $2 - AND mcc.id = $5 - AND mcc.runtime_lock_id = $3 -), -inserted_call AS ( - INSERT INTO tool_calls( - id, agent_id, model_output_id, - provider_call_id, name, input, - type, state, created_at - ) - SELECT id, agent_id, model_output_id, - provider_call_id, name, input, - type, 'awaiting_authorization', created_at - FROM candidate - RETURNING id, agent_id, - model_output_id, provider_call_id, name, - input, type, state, - runtime_lock_id, created_at -) -SELECT call.id, candidate.project_id, call.agent_id, - candidate.turn_id, candidate.source_event_id, - candidate.model_call_context_id, call.provider_call_id, - call.name, call.input, - call.type, call.state, - ''::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at, NULL::timestamptz AS completed_at -FROM inserted_call call -JOIN candidate ON candidate.id = call.id + SELECT coalesce(jsonb_agg( + CASE + WHEN block.block_kind = 'text' THEN jsonb_build_object('type', 'text', 'text', block.text_content) + WHEN block.block_kind = 'structured_data' THEN jsonb_build_object('type', 'structured_data', 'value', block.structured_data) + WHEN block.block_kind = 'artifact' THEN + jsonb_build_object('type', 'media_ref', 'artifact_id', block.artifact_id::text) || + CASE + WHEN block.exclude_from_model_context THEN jsonb_build_object('exclude_from_model_context', true) + ELSE '{}'::jsonb + END + END + ORDER BY block.ordinal, block.id + ) FILTER (WHERE block.id IS NOT NULL AND block.block_kind IN ('text', 'structured_data', 'artifact')), '[]'::jsonb) AS content_parts + FROM content_blocks block + WHERE block.agent_id = result.agent_id + AND block.owner_tool_call_result_id = result.id +) result_blocks ON true ` -type InsertToolCallParams struct { +type GetToolCallByProviderCallParams struct { ProjectID uuid.UUID AgentID uuid.UUID - RuntimeLockID uuid.UUID - SourceEventID uuid.UUID ModelCallContextID uuid.UUID - ToolCallID *uuid.UUID ProviderCallID string - Name string - Input json.RawMessage - Type string } -type InsertToolCallRow struct { +type GetToolCallByProviderCallRow struct { ID uuid.UUID ProjectID uuid.UUID AgentID uuid.UUID @@ -1216,20 +171,14 @@ type InsertToolCallRow struct { CompletedAt *time.Time } -func (q *Queries) InsertToolCall(ctx context.Context, arg InsertToolCallParams) (InsertToolCallRow, error) { - row := q.db.QueryRow(ctx, insertToolCall, +func (q *Queries) GetToolCallByProviderCall(ctx context.Context, arg GetToolCallByProviderCallParams) (GetToolCallByProviderCallRow, error) { + row := q.db.QueryRow(ctx, getToolCallByProviderCall, arg.ProjectID, arg.AgentID, - arg.RuntimeLockID, - arg.SourceEventID, arg.ModelCallContextID, - arg.ToolCallID, arg.ProviderCallID, - arg.Name, - arg.Input, - arg.Type, ) - var i InsertToolCallRow + var i GetToolCallByProviderCallRow err := row.Scan( &i.ID, &i.ProjectID, @@ -1252,6 +201,17 @@ func (q *Queries) InsertToolCall(ctx context.Context, arg InsertToolCallParams) } const listCompletedToolCallsAtWatermark = `-- name: ListCompletedToolCallsAtWatermark :many +WITH selected_calls AS MATERIALIZED ( + SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, + tc.model_call_context_id, tc.model_output_id, tc.provider_call_id, tc.name, tc.input, + tc.type, tc.state, tc.runtime_lock_id, tc.created_at + FROM tool_call_read_projection tc + WHERE tc.project_id = $3 + AND tc.agent_id = $4 + AND tc.state = 'completed' + AND tc.source_event_sequence > $1 + AND tc.source_event_sequence <= $2 +) SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, tc.model_call_context_id, result.id AS tool_call_result_id, @@ -1261,7 +221,7 @@ SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.runtime_lock_id, coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, tc.created_at, result.completed_at -FROM tool_call_read_projection tc +FROM selected_calls tc JOIN content_blocks call_block ON call_block.agent_id = tc.agent_id AND call_block.tool_call_id = tc.id @@ -1291,21 +251,16 @@ LEFT JOIN LATERAL ( WHERE block.agent_id = result.agent_id AND block.owner_tool_call_result_id = result.id ) result_blocks ON true -WHERE tc.project_id = $1 - AND tc.agent_id = $2 - AND tc.state = 'completed' - AND tc.source_event_sequence > $3 - AND tool_result_event.sequence > $3 - AND tc.source_event_sequence <= $4 - AND tool_result_event.sequence <= $4 +WHERE tool_result_event.sequence > $1 + AND tool_result_event.sequence <= $2 ORDER BY tc.source_event_sequence, call_block.ordinal, tc.created_at, tc.id ` type ListCompletedToolCallsAtWatermarkParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID AfterEventSequence int64 MaxEventSequence int64 + ProjectID uuid.UUID + AgentID uuid.UUID } type ListCompletedToolCallsAtWatermarkRow struct { @@ -1333,10 +288,10 @@ type ListCompletedToolCallsAtWatermarkRow struct { func (q *Queries) ListCompletedToolCallsAtWatermark(ctx context.Context, arg ListCompletedToolCallsAtWatermarkParams) ([]ListCompletedToolCallsAtWatermarkRow, error) { rows, err := q.db.Query(ctx, listCompletedToolCallsAtWatermark, - arg.ProjectID, - arg.AgentID, arg.AfterEventSequence, arg.MaxEventSequence, + arg.ProjectID, + arg.AgentID, ) if err != nil { return nil, err @@ -1478,319 +433,6 @@ func (q *Queries) ListToolCallsForAgent(ctx context.Context, arg ListToolCallsFo return items, nil } -const listToolCallsForModelContext = `-- name: ListToolCallsForModelContext :many -SELECT tc.id, tc.project_id, tc.agent_id, - tc.turn_id, tc.source_event_id, tc.model_call_context_id, - tc.provider_call_id, tc.name, tc.input, tc.type, tc.state, - coalesce(result.outcome, '') AS outcome, tc.runtime_lock_id, - coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, - tc.created_at, result.completed_at -FROM tool_call_read_projection tc -LEFT JOIN tool_call_results result - ON result.agent_id = tc.agent_id - AND result.tool_call_id = tc.id -LEFT JOIN LATERAL ( - SELECT coalesce(jsonb_agg( - CASE - WHEN block.block_kind = 'text' THEN jsonb_build_object('type', 'text', 'text', block.text_content) - WHEN block.block_kind = 'structured_data' THEN jsonb_build_object('type', 'structured_data', 'value', block.structured_data) - WHEN block.block_kind = 'artifact' THEN - jsonb_build_object('type', 'media_ref', 'artifact_id', block.artifact_id::text) - || CASE - WHEN block.exclude_from_model_context THEN jsonb_build_object('exclude_from_model_context', true) - ELSE '{}'::jsonb - END - END - ORDER BY block.ordinal, block.id - ) FILTER ( - WHERE block.id IS NOT NULL - AND block.block_kind IN ('text', 'structured_data', 'artifact') - ), '[]'::jsonb) AS content_parts - FROM content_blocks block - WHERE block.agent_id = result.agent_id - AND block.owner_tool_call_result_id = result.id -) result_blocks ON true -WHERE tc.project_id = $1 - AND tc.agent_id = $2 - AND tc.model_call_context_id = $3 -ORDER BY tc.created_at, tc.id -` - -type ListToolCallsForModelContextParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - ModelCallContextID uuid.UUID -} - -type ListToolCallsForModelContextRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time - CompletedAt *time.Time -} - -func (q *Queries) ListToolCallsForModelContext(ctx context.Context, arg ListToolCallsForModelContextParams) ([]ListToolCallsForModelContextRow, error) { - rows, err := q.db.Query(ctx, listToolCallsForModelContext, arg.ProjectID, arg.AgentID, arg.ModelCallContextID) - if err != nil { - return nil, err - } - defer rows.Close() - items := []ListToolCallsForModelContextRow{} - for rows.Next() { - var i ListToolCallsForModelContextRow - if err := rows.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - &i.CompletedAt, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - -const markToolCallAwaitingPermission = `-- name: MarkToolCallAwaitingPermission :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $3 - AND runtime_lock.agent_id = $4 - AND runtime_lock.id = $2 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'awaiting_permission' -FROM live_runtime runtime_lock -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = $1 - AND call.state = 'awaiting_authorization' - AND runtime_lock.id = $2 -RETURNING call.id -` - -type MarkToolCallAwaitingPermissionParams struct { - ID uuid.UUID - RuntimeLockID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -func (q *Queries) MarkToolCallAwaitingPermission(ctx context.Context, arg MarkToolCallAwaitingPermissionParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, markToolCallAwaitingPermission, - arg.ID, - arg.RuntimeLockID, - arg.ProjectID, - arg.AgentID, - ) - var id uuid.UUID - err := row.Scan(&id) - return id, err -} - -const markToolCallReady = `-- name: MarkToolCallReady :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $3 - AND runtime_lock.agent_id = $4 - AND runtime_lock.id = $2 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'ready' -FROM live_runtime runtime_lock -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = $1 - AND call.state = 'awaiting_authorization' - AND runtime_lock.id = $2 - AND projection.project_id = runtime_lock.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - ''::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at, NULL::timestamptz AS completed_at -` - -type MarkToolCallReadyParams struct { - ID uuid.UUID - RuntimeLockID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID -} - -type MarkToolCallReadyRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time - CompletedAt *time.Time -} - -func (q *Queries) MarkToolCallReady(ctx context.Context, arg MarkToolCallReadyParams) (MarkToolCallReadyRow, error) { - row := q.db.QueryRow(ctx, markToolCallReady, - arg.ID, - arg.RuntimeLockID, - arg.ProjectID, - arg.AgentID, - ) - var i MarkToolCallReadyRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - &i.CompletedAt, - ) - return i, err -} - -const markToolCallReadyFromInteraction = `-- name: MarkToolCallReadyFromInteraction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = $1 - AND agent.id = $2 - FOR UPDATE -), -permission_interaction AS MATERIALIZED ( - SELECT agent.project_id, interaction.agent_id, interaction.tool_call_id - FROM agent_interactions interaction - JOIN locked_agent agent ON agent.id = interaction.agent_id - WHERE interaction.agent_id = $2 - AND interaction.id = $3 - AND interaction.interaction_kind = 'permission' - AND interaction.state = 'resolved' -) -UPDATE tool_calls call -SET state = 'ready' -FROM permission_interaction interaction -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = interaction.agent_id - AND call.id = interaction.tool_call_id - AND call.state = 'awaiting_permission' - AND projection.project_id = interaction.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - ''::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at, NULL::timestamptz AS completed_at -` - -type MarkToolCallReadyFromInteractionParams struct { - ProjectID uuid.UUID - AgentID uuid.UUID - InteractionID uuid.UUID -} - -type MarkToolCallReadyFromInteractionRow struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - TurnID uuid.UUID - SourceEventID uuid.UUID - ModelCallContextID uuid.UUID - ProviderCallID string - Name string - Input json.RawMessage - Type string - State string - Outcome string - RuntimeLockID *uuid.UUID - ResultContentParts json.RawMessage - CreatedAt time.Time - CompletedAt *time.Time -} - -func (q *Queries) MarkToolCallReadyFromInteraction(ctx context.Context, arg MarkToolCallReadyFromInteractionParams) (MarkToolCallReadyFromInteractionRow, error) { - row := q.db.QueryRow(ctx, markToolCallReadyFromInteraction, arg.ProjectID, arg.AgentID, arg.InteractionID) - var i MarkToolCallReadyFromInteractionRow - err := row.Scan( - &i.ID, - &i.ProjectID, - &i.AgentID, - &i.TurnID, - &i.SourceEventID, - &i.ModelCallContextID, - &i.ProviderCallID, - &i.Name, - &i.Input, - &i.Type, - &i.State, - &i.Outcome, - &i.RuntimeLockID, - &i.ResultContentParts, - &i.CreatedAt, - &i.CompletedAt, - ) - return i, err -} - const nextRunnableToolCallForModelOutput = `-- name: NextRunnableToolCallForModelOutput :one SELECT call.id, call.project_id, call.agent_id, call.turn_id, call.source_event_id, @@ -1873,100 +515,3 @@ func (q *Queries) NextRunnableToolCallForModelOutput(ctx context.Context, arg Ne ) return i, err } - -const requeueRuntimeToolCall = `-- name: RequeueRuntimeToolCall :execrows -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $2 - AND runtime_lock.agent_id = $3 - AND runtime_lock.id = $4 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'ready', - runtime_lock_id = NULL -FROM live_runtime runtime_lock -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = $1 - AND call.state = 'running' - AND call.runtime_lock_id = runtime_lock.id -` - -type RequeueRuntimeToolCallParams struct { - ID uuid.UUID - ProjectID uuid.UUID - AgentID uuid.UUID - RuntimeLockID uuid.UUID -} - -func (q *Queries) RequeueRuntimeToolCall(ctx context.Context, arg RequeueRuntimeToolCallParams) (int64, error) { - result, err := q.db.Exec(ctx, requeueRuntimeToolCall, - arg.ID, - arg.ProjectID, - arg.AgentID, - arg.RuntimeLockID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} - -const startToolCall = `-- name: StartToolCall :execrows -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = $5 - AND runtime_lock.agent_id = $2 - AND runtime_lock.id = $4 - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = CASE - WHEN $1::boolean THEN 'running' - ELSE 'waiting' - END, - runtime_lock_id = CASE - WHEN $1::boolean THEN runtime_lock.id - ELSE NULL - END -FROM live_runtime runtime_lock -WHERE call.agent_id = $2 - AND call.id = $3 - AND call.state = 'ready' - AND call.runtime_lock_id IS NULL - AND call.type IN ('built_in', 'mcp') - AND ( - $1::boolean - OR call.type = 'built_in' - ) - AND runtime_lock.agent_id = call.agent_id - AND runtime_lock.id = $4 -` - -type StartToolCallParams struct { - RetainRuntimeOwnership bool - AgentID uuid.UUID - ID uuid.UUID - RuntimeLockID uuid.UUID - ProjectID uuid.UUID -} - -func (q *Queries) StartToolCall(ctx context.Context, arg StartToolCallParams) (int64, error) { - result, err := q.db.Exec(ctx, startToolCall, - arg.RetainRuntimeOwnership, - arg.AgentID, - arg.ID, - arg.RuntimeLockID, - arg.ProjectID, - ) - if err != nil { - return 0, err - } - return result.RowsAffected(), nil -} diff --git a/internal/storage/internal/lifecyclelock/lifecyclelock.go b/internal/storage/internal/lifecyclelock/lifecyclelock.go index 20e400d45..1fe99aa33 100644 --- a/internal/storage/internal/lifecyclelock/lifecyclelock.go +++ b/internal/storage/internal/lifecyclelock/lifecyclelock.go @@ -42,7 +42,7 @@ type AgentRef struct { AgentID uuid.UUID } -func OrganizationShared(ctx context.Context, tx pgx.Tx, orgID uuid.UUID) error { +func OrganizationShared(ctx context.Context, tx dbsqlc.DBTX, orgID uuid.UUID) error { return organizationShared(ctx, dbsqlc.New(tx), orgID) } @@ -56,7 +56,7 @@ func organizationShared(ctx context.Context, q *dbsqlc.Queries, orgID uuid.UUID) return nil } -func EnterActiveOrganization(ctx context.Context, tx pgx.Tx, orgID uuid.UUID) error { +func EnterActiveOrganization(ctx context.Context, tx dbsqlc.DBTX, orgID uuid.UUID) error { q := dbsqlc.New(tx) return enterActiveOrganization(ctx, q, orgID) } @@ -76,7 +76,7 @@ func enterActiveOrganization(ctx context.Context, q *dbsqlc.Queries, orgID uuid. return nil } -func OrganizationExclusive(ctx context.Context, tx pgx.Tx, orgID uuid.UUID) error { +func OrganizationExclusive(ctx context.Context, tx dbsqlc.DBTX, orgID uuid.UUID) error { q := dbsqlc.New(tx) if err := q.LockOrganizationLifecycleExclusive( ctx, @@ -99,7 +99,7 @@ func projectShared(ctx context.Context, q *dbsqlc.Queries, projectID uuid.UUID) func EnterActiveProject( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, orgID, projectID uuid.UUID, ) error { return EnterActiveProjects(ctx, tx, orgID, []uuid.UUID{projectID}) @@ -107,7 +107,7 @@ func EnterActiveProject( func EnterActiveProjects( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, orgID uuid.UUID, projectIDs []uuid.UUID, ) error { @@ -138,7 +138,7 @@ func EnterActiveProjects( return nil } -func ProjectExclusive(ctx context.Context, tx pgx.Tx, projectID uuid.UUID) error { +func ProjectExclusive(ctx context.Context, tx dbsqlc.DBTX, projectID uuid.UUID) error { if err := dbsqlc.New(tx).LockProjectLifecycleExclusive( ctx, dbsqlc.LockProjectLifecycleExclusiveParams{ProjectID: projectID}, @@ -165,7 +165,7 @@ func orderedIDs(ids []uuid.UUID) []uuid.UUID { return deduped } -func AgentSources(ctx context.Context, tx pgx.Tx, agentIDs ...uuid.UUID) error { +func AgentSources(ctx context.Context, tx dbsqlc.DBTX, agentIDs ...uuid.UUID) error { q := dbsqlc.New(tx) for _, agentID := range orderedIDs(agentIDs) { if err := q.LockAgentMachineSources( @@ -178,7 +178,7 @@ func AgentSources(ctx context.Context, tx pgx.Tx, agentIDs ...uuid.UUID) error { return nil } -func Pools(ctx context.Context, tx pgx.Tx, refs []PoolRef) error { +func Pools(ctx context.Context, tx dbsqlc.DBTX, refs []PoolRef) error { q := dbsqlc.New(tx) ordered := append([]PoolRef(nil), refs...) sort.Slice(ordered, func(i, j int) bool { @@ -203,7 +203,7 @@ func Pools(ctx context.Context, tx pgx.Tx, refs []PoolRef) error { return nil } -func PoolGrants(ctx context.Context, tx pgx.Tx, grantIDs []uuid.UUID) error { +func PoolGrants(ctx context.Context, tx dbsqlc.DBTX, grantIDs []uuid.UUID) error { q := dbsqlc.New(tx) for _, grantID := range orderedIDs(grantIDs) { if _, err := q.LockProjectMachinePoolGrantForLifecycle( @@ -218,7 +218,7 @@ func PoolGrants(ctx context.Context, tx pgx.Tx, grantIDs []uuid.UUID) error { return nil } -func Machines(ctx context.Context, tx pgx.Tx, refs []MachineRef) error { +func Machines(ctx context.Context, tx dbsqlc.DBTX, refs []MachineRef) error { q := dbsqlc.New(tx) ordered := append([]MachineRef(nil), refs...) sort.Slice(ordered, func(i, j int) bool { @@ -243,7 +243,7 @@ func Machines(ctx context.Context, tx pgx.Tx, refs []MachineRef) error { return nil } -func Agents(ctx context.Context, tx pgx.Tx, refs []AgentRef) error { +func Agents(ctx context.Context, tx dbsqlc.DBTX, refs []AgentRef) error { q := dbsqlc.New(tx) for _, ref := range orderedAgentRefs(refs) { if _, err := q.LockAgentInProject( @@ -256,22 +256,6 @@ func Agents(ctx context.Context, tx pgx.Tx, refs []AgentRef) error { return nil } -func TryAgents(ctx context.Context, tx pgx.Tx, refs []AgentRef) (bool, error) { - q := dbsqlc.New(tx) - for _, ref := range orderedAgentRefs(refs) { - if _, err := q.TryLockAgentInProject( - ctx, - dbsqlc.TryLockAgentInProjectParams{ProjectID: ref.ProjectID, ID: ref.AgentID}, - ); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return false, nil - } - return false, fmt.Errorf("try lock agent for lifecycle: %w", err) - } - } - return true, nil -} - func rowLockError(operation string, err error) error { if errors.Is(err, pgx.ErrNoRows) { return storeerr.ErrNotFound diff --git a/internal/storage/internal/secretops/secretops.go b/internal/storage/internal/secretops/secretops.go index a06604119..c367e33bc 100644 --- a/internal/storage/internal/secretops/secretops.go +++ b/internal/storage/internal/secretops/secretops.go @@ -5,7 +5,6 @@ import ( "fmt" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/secrets" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -24,7 +23,7 @@ type Facts struct { func LockReference( ctx context.Context, - tx pgx.Tx, + tx dbsqlc.DBTX, orgID, secretID uuid.UUID, ) (Facts, error) { q := dbsqlc.New(tx) diff --git a/internal/storage/internal/storeutil/transactions.go b/internal/storage/internal/storeutil/transactions.go index 5df9e7d12..1d8c978eb 100644 --- a/internal/storage/internal/storeutil/transactions.go +++ b/internal/storage/internal/storeutil/transactions.go @@ -3,11 +3,9 @@ package storeutil import ( "context" "errors" - "fmt" "time" "github.com/jackc/pgx/v5/pgconn" - "github.com/omnara-ai/omnara/internal/notifications" log "github.com/omnara-ai/omnara/observability/wideevent" ) @@ -89,17 +87,3 @@ func retryableTransactionSQLState(err error) string { } return "" } - -func CommitTxWithNotifications( - ctx context.Context, - tx interface{ Commit(context.Context) error }, - txNotifications *notifications.TxNotifications, - publisher notifications.PostCommitPublisher, - operation string, -) error { - if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("commit %s: %w", operation, err) - } - txNotifications.Flush(context.WithoutCancel(ctx), publisher) - return nil -} diff --git a/internal/storage/internal/storeutil/transactions_test.go b/internal/storage/internal/storeutil/transactions_test.go index 85082cfb1..271477a59 100644 --- a/internal/storage/internal/storeutil/transactions_test.go +++ b/internal/storage/internal/storeutil/transactions_test.go @@ -7,23 +7,14 @@ import ( "errors" "fmt" "log/slog" - "reflect" "strings" "testing" - "github.com/google/uuid" "github.com/jackc/pgx/v5/pgconn" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" log "github.com/omnara-ai/omnara/observability/wideevent" ) -type recordingCommitter struct { - err error - order *[]string - afterRun func() -} - func TestRetryTransactionRetriesPostgresTransactionConflicts(t *testing.T) { for _, code := range []string{"40P01", "40001"} { t.Run(code, func(t *testing.T) { @@ -86,90 +77,6 @@ func TestRetryTransactionHonorsContextCancellation(t *testing.T) { } } -func (c recordingCommitter) Commit(context.Context) error { - *c.order = append(*c.order, "commit") - if c.afterRun != nil { - c.afterRun() - } - return c.err -} - -type recordingPublisher struct { - order *[]string - contextErr error - intents []notifications.PostCommitIntent -} - -func (p *recordingPublisher) PublishPostCommit( - ctx context.Context, - intent notifications.PostCommitIntent, -) { - *p.order = append(*p.order, "publish") - p.contextErr = ctx.Err() - p.intents = append(p.intents, intent) -} - -func TestCommitTxWithNotificationsCommitsBeforePublishing(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - order := []string{} - agentID := uuid.New() - txNotifications := notifications.NewTxNotifications() - txNotifications.AddAgentEvent(agentID, 1, "agent_input") - publisher := &recordingPublisher{order: &order} - - err := storeutil.CommitTxWithNotifications( - ctx, - recordingCommitter{order: &order, afterRun: cancel}, - txNotifications, - publisher, - "test operation", - ) - if err != nil { - t.Fatalf("CommitTxWithNotifications() error = %v", err) - } - if !reflect.DeepEqual(order, []string{"commit", "publish"}) { - t.Fatalf("operation order = %v", order) - } - if publisher.contextErr != nil { - t.Fatalf("publish context error = %v", publisher.contextErr) - } - if len(publisher.intents) != 1 { - t.Fatalf("published intents = %d", len(publisher.intents)) - } - intent, ok := publisher.intents[0].(notifications.AgentEventCommitted) - if !ok || intent.AgentID != agentID { - t.Fatalf("published intent = %#v", publisher.intents[0]) - } -} - -func TestCommitTxWithNotificationsDoesNotPublishAfterCommitFailure(t *testing.T) { - commitErr := errors.New("commit failed") - order := []string{} - txNotifications := notifications.NewTxNotifications() - txNotifications.AddAgentEvent(uuid.New(), 1, "agent_input") - publisher := &recordingPublisher{order: &order} - - err := storeutil.CommitTxWithNotifications( - context.Background(), - recordingCommitter{err: commitErr, order: &order}, - txNotifications, - publisher, - "test operation", - ) - if !errors.Is(err, commitErr) { - t.Fatalf("CommitTxWithNotifications() error = %v", err) - } - if err.Error() != "commit test operation: commit failed" { - t.Fatalf("CommitTxWithNotifications() error = %q", err) - } - if !reflect.DeepEqual(order, []string{"commit"}) { - t.Fatalf("operation order = %v", order) - } - if len(publisher.intents) != 0 { - t.Fatalf("published intents = %d", len(publisher.intents)) - } -} - func TestRetryTransactionDoesNotLogSuccessOnPanic(t *testing.T) { t.Parallel() var output bytes.Buffer diff --git a/internal/storage/orglifecycle/orgs_projects.go b/internal/storage/orglifecycle/orgs_projects.go index b5a9243ee..a9e911de7 100644 --- a/internal/storage/orglifecycle/orgs_projects.go +++ b/internal/storage/orglifecycle/orgs_projects.go @@ -7,10 +7,10 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" - "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/resourcename" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/skillops" @@ -77,7 +77,7 @@ func (s *Service) CreateOrgForUser( return identitystore.CreateOrgForUserRecord{}, err } if input.ProvisionDefaultModelProvider { - rows, err := s.q.WithTx(tx).EnqueueDefaultModelProviderProvisioning( + rows, err := dbsqlc.New(tx).EnqueueDefaultModelProviderProvisioning( ctx, dbsqlc.EnqueueDefaultModelProviderProvisioningParams{ OrganizationID: record.Org.ID, @@ -190,15 +190,14 @@ func deleteProjectOwnedContentTx( func (s *Service) teardownProjectAgentsTx( ctx context.Context, - tx pgx.Tx, - txNotifications *notifications.TxNotifications, + unit *agentexecution.Unit, projectID uuid.UUID, agentIDs []uuid.UUID, actor *executionstore.ActorParams, ) ([]executionstore.MachineRecord, error) { machines := make([]executionstore.MachineRecord, 0) for _, agentID := range agentIDs { - agentMachines, err := s.execution.ArchiveAgentTx(ctx, tx, txNotifications, projectID, agentID, actor) + agentMachines, err := s.execution.ArchiveAgentInUnit(ctx, unit, projectID, agentID, actor) if err != nil { return nil, fmt.Errorf("archive project agent %s: %w", agentID, err) } @@ -215,10 +214,11 @@ type organizationMachineLifecyclePlan struct { func prelockProjectMachineLifecycleTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, q *dbsqlc.Queries, orgID, projectID uuid.UUID, ) ([]uuid.UUID, error) { + tx := unit.DB() agentIDs, err := q.ListActiveAgentIDsForProjectDeletion( ctx, dbsqlc.ListActiveAgentIDsForProjectDeletionParams{ProjectID: projectID}, @@ -276,7 +276,7 @@ func prelockProjectMachineLifecycleTx( if err := lifecyclelock.Machines(ctx, tx, machineRefs); err != nil { return nil, err } - if err := lifecyclelock.Agents(ctx, tx, agentRefs); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { return nil, err } return agentIDs, nil @@ -284,10 +284,11 @@ func prelockProjectMachineLifecycleTx( func prelockOrganizationMachineLifecycleTx( ctx context.Context, - tx pgx.Tx, + unit *agentexecution.Unit, q *dbsqlc.Queries, orgID uuid.UUID, ) (organizationMachineLifecyclePlan, error) { + tx := unit.DB() agentRows, err := q.ListActiveAgentRefsForOrganizationDeletion( ctx, dbsqlc.ListActiveAgentRefsForOrganizationDeletionParams{OrgID: orgID}, @@ -373,7 +374,7 @@ func prelockOrganizationMachineLifecycleTx( if err := lifecyclelock.Machines(ctx, tx, machineRefs); err != nil { return organizationMachineLifecyclePlan{}, err } - if err := lifecyclelock.Agents(ctx, tx, agentRefs); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { return organizationMachineLifecyclePlan{}, err } byoMachineIDs, err := q.ListActiveBYOMachineIDsForOrganizationDeletion( @@ -430,16 +431,15 @@ func (s *Service) deleteProjectOnce( orgID, projectID uuid.UUID, actor *executionstore.ActorParams, ) ([]executionstore.MachineRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return nil, fmt.Errorf("begin delete project: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + q := dbsqlc.New(tx) - if err := lifecyclelock.OrganizationShared(ctx, tx, orgID); err != nil { - return nil, err - } - if err := lifecyclelock.ProjectExclusive(ctx, tx, projectID); err != nil { + if err := unit.LockProjectDeletion(ctx, orgID, projectID); err != nil { return nil, err } if _, err := q.GetProject(ctx, dbsqlc.GetProjectParams{OrgID: orgID, ID: projectID}); err != nil { @@ -448,15 +448,13 @@ func (s *Service) deleteProjectOnce( } return nil, fmt.Errorf("load project for deletion: %w", err) } - agentIDs, err := prelockProjectMachineLifecycleTx(ctx, tx, q, orgID, projectID) + agentIDs, err := prelockProjectMachineLifecycleTx(ctx, unit, q, orgID, projectID) if err != nil { return nil, err } - txNotifications := s.newTxNotifications() machines, err := s.teardownProjectAgentsTx( ctx, - tx, - txNotifications, + unit, projectID, agentIDs, actor, @@ -476,6 +474,9 @@ func (s *Service) deleteProjectOnce( if err != nil { return nil, err } + if err := unit.CaptureWebhookEligibility(ctx); err != nil { + return nil, err + } rows, err := q.DeleteProject(ctx, dbsqlc.DeleteProjectParams{OrgID: orgID, ID: projectID}) if err != nil { return nil, fmt.Errorf("delete project: %w", err) @@ -483,13 +484,7 @@ func (s *Service) deleteProjectOnce( if rows == 0 { return nil, storeerr.ErrNotFound } - if err := storeutil.CommitTxWithNotifications( - ctx, - tx, - txNotifications, - s.postCommitPublisher, - "delete project", - ); err != nil { + if err := unit.Commit(ctx, "delete project"); err != nil { return nil, err } skillops.Purge(ctx, s.blobs, skillArchives) @@ -522,13 +517,15 @@ func (s *Service) deleteOrganizationOnce( orgID uuid.UUID, actor *executionstore.ActorParams, ) ([]executionstore.MachineRecord, error) { - tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{}) + unit, err := s.cell.Begin(ctx) if err != nil { return nil, fmt.Errorf("begin delete organization: %w", err) } - defer func() { _ = tx.Rollback(ctx) }() + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + q := dbsqlc.New(tx) - if err := lifecyclelock.OrganizationExclusive(ctx, tx, orgID); err != nil { + if err := unit.LockOrganizationDeletion(ctx, orgID); err != nil { return nil, err } active, err := q.OrgExistsActive(ctx, dbsqlc.OrgExistsActiveParams{ID: orgID}) @@ -545,11 +542,17 @@ func (s *Service) deleteOrganizationOnce( if err != nil { return nil, fmt.Errorf("list organization projects: %w", err) } - plan, err := prelockOrganizationMachineLifecycleTx(ctx, tx, q, orgID) + plan, err := prelockOrganizationMachineLifecycleTx(ctx, unit, q, orgID) if err != nil { return nil, err } - txNotifications := s.newTxNotifications() + var webhookAgents []uuid.UUID + for _, ids := range plan.agentIDsByProject { + webhookAgents = append(webhookAgents, ids...) + } + if err := unit.CaptureWebhookEligibility(ctx, webhookAgents...); err != nil { + return nil, err + } rows, err := q.DeleteOrganization(ctx, dbsqlc.DeleteOrganizationParams{ID: orgID}) if err != nil { return nil, fmt.Errorf("delete organization: %w", err) @@ -561,8 +564,7 @@ func (s *Service) deleteOrganizationOnce( for _, projectID := range orgProjectIDs { projectMachines, err := s.teardownProjectAgentsTx( ctx, - tx, - txNotifications, + unit, projectID, plan.agentIDsByProject[projectID], actor, @@ -573,14 +575,14 @@ func (s *Service) deleteOrganizationOnce( machines = append(machines, projectMachines...) } for _, poolID := range plan.poolIDs { - poolMachines, err := s.execution.DeleteMachinePoolTx(ctx, tx, txNotifications, orgID, poolID) + poolMachines, err := s.execution.DeleteMachinePoolInUnit(ctx, unit, orgID, poolID) if err != nil { return nil, fmt.Errorf("delete organization machine pool %s: %w", poolID, err) } machines = append(machines, poolMachines...) } for _, machineID := range plan.byoMachineIDs { - _, err := s.execution.DeleteMachineTx(ctx, tx, txNotifications, executionstore.DeleteMachineInput{ + _, err := s.execution.DeleteMachineInUnit(ctx, unit, executionstore.DeleteMachineInput{ OrgID: orgID, MachineID: machineID, }) if err != nil { @@ -652,13 +654,7 @@ func (s *Service) deleteOrganizationOnce( if err := q.DeleteOrganizationOrgAPIKeys(ctx, dbsqlc.DeleteOrganizationOrgAPIKeysParams{OrgID: orgID}); err != nil { return nil, fmt.Errorf("delete organization api keys: %w", err) } - if err := storeutil.CommitTxWithNotifications( - ctx, - tx, - txNotifications, - s.postCommitPublisher, - "delete organization", - ); err != nil { + if err := unit.Commit(ctx, "delete organization"); err != nil { return nil, err } skillops.Purge(ctx, s.blobs, skillArchives) diff --git a/internal/storage/orglifecycle/service.go b/internal/storage/orglifecycle/service.go index 702ce8595..70c0aff71 100644 --- a/internal/storage/orglifecycle/service.go +++ b/internal/storage/orglifecycle/service.go @@ -6,6 +6,7 @@ import ( "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/modelstore" @@ -22,30 +23,26 @@ type Config struct { } type Service struct { - pool *storeutil.Pool - q *dbsqlc.Queries - blobs blobstore.Store - postCommitPublisher notifications.PostCommitPublisher - identity *identitystore.Store - execution *executionstore.Store - models *modelstore.Store - secrets *secretstore.Store + cell *agentexecution.AgentCell + pool *storeutil.Pool + q *dbsqlc.Queries + blobs blobstore.Store + identity *identitystore.Store + execution *executionstore.Store + models *modelstore.Store + secrets *secretstore.Store } func New(pool *pgxpool.Pool, config Config) *Service { db := storeutil.WrapPool(pool) return &Service{ - pool: db, - q: dbsqlc.New(db), - blobs: config.Blobs, - postCommitPublisher: config.PostCommitPublisher, - identity: config.Identity, - execution: config.Execution, - models: config.Models, - secrets: config.Secrets, + cell: agentexecution.NewCell("primary", pool, config.PostCommitPublisher, nil), + pool: db, + q: dbsqlc.New(db), + blobs: config.Blobs, + identity: config.Identity, + execution: config.Execution, + models: config.Models, + secrets: config.Secrets, } } - -func (s *Service) newTxNotifications() *notifications.TxNotifications { - return notifications.NewTxNotifications() -} diff --git a/internal/storage/queries/actors.sql b/internal/storage/queries/actors.sql index 71ef136c2..9b0c111e4 100644 --- a/internal/storage/queries/actors.sql +++ b/internal/storage/queries/actors.sql @@ -42,12 +42,6 @@ FROM actors WHERE project_id = sqlc.arg(project_id) AND id = sqlc.arg(id); --- name: ListActorIdentitiesByIDs :many -SELECT id, provider, provider_user_id -FROM actors -WHERE project_id = sqlc.arg(project_id) - AND id = ANY(sqlc.arg(ids)::uuid[]); - -- name: ListActors :many SELECT id, project_id, provider, provider_tenant_id, provider_user_id, display_name, metadata, created_at, updated_at FROM actors @@ -80,3 +74,9 @@ WHERE project_id = sqlc.arg(project_id) AND provider_tenant_id IS NOT DISTINCT FROM sqlc.narg(provider_tenant_id) AND provider_user_id = sqlc.arg(provider_user_id) AND display_name IS DISTINCT FROM sqlc.arg(display_name)::text; + +-- name: ListActorIdentitiesByIDs :many +SELECT id, provider, provider_user_id +FROM actors +WHERE project_id = sqlc.arg(project_id) + AND id = ANY(sqlc.arg(ids)::uuid[]); diff --git a/internal/storage/queries/agent_inputs.sql b/internal/storage/queries/agent_inputs.sql index a4acb202b..8cdac3ce6 100644 --- a/internal/storage/queries/agent_inputs.sql +++ b/internal/storage/queries/agent_inputs.sql @@ -1,58 +1,3 @@ --- name: MarkAgentWakeup :exec -INSERT INTO agent_wakeups(agent_id, ready_at, updated_at, metadata) -SELECT agent.id, - statement_timestamp(), statement_timestamp(), sqlc.arg(metadata) -FROM agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - AND agent.state <> 'archived' -ON CONFLICT (agent_id) DO UPDATE -SET ready_at = LEAST(agent_wakeups.ready_at, excluded.ready_at), - updated_at = GREATEST(agent_wakeups.updated_at, excluded.updated_at), - metadata = agent_wakeups.metadata || excluded.metadata; - --- name: ConsumeAgentWakeup :execrows -DELETE FROM agent_wakeups wake -USING agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - AND wake.agent_id = agent.id; - --- name: ReconcileAgentWakeup :exec -WITH locked_agent AS MATERIALIZED ( - SELECT agent.id, agent.project_id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - AND agent.state <> 'archived' -), -desired AS MATERIALIZED ( - SELECT agent_next_wakeup_ready_at( - locked_agent.project_id, - locked_agent.id - ) AS ready_at - FROM locked_agent -), -upserted AS ( - INSERT INTO agent_wakeups(agent_id, ready_at, updated_at, metadata) - SELECT locked_agent.id, - desired.ready_at, - statement_timestamp(), - sqlc.arg(metadata) - FROM locked_agent - JOIN desired ON desired.ready_at IS NOT NULL - ON CONFLICT (agent_id) DO UPDATE - SET ready_at = excluded.ready_at, - updated_at = GREATEST(agent_wakeups.updated_at, excluded.updated_at), - metadata = agent_wakeups.metadata || excluded.metadata - RETURNING agent_id -) -DELETE FROM agent_wakeups wake -USING locked_agent -CROSS JOIN desired -WHERE wake.agent_id = locked_agent.id - AND desired.ready_at IS NULL; - -- name: ClaimNextAgentWakeup :one -- The query walks wakeups in global ready order, but locks agents before wake -- rows to match all other agent mutation paths. If the selected wake is no @@ -87,42 +32,6 @@ locked_wake AS MATERIALIZED ( SELECT agent_id, project_id FROM locked_wake; --- name: DeleteAgentWakeup :execrows -WITH locked AS ( - SELECT agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -DELETE FROM agent_wakeups wake -USING locked -WHERE wake.agent_id = locked.id; - --- name: ListSteeringAgentInputsForAdmission :many -SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata -FROM agent_inputs input -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.state = 'received' - AND input.delivery_mode = 'steering' - AND input.input_kind = 'content' - AND NOT agent_has_incomplete_tool_batch(sqlc.arg(project_id), sqlc.arg(agent_id)) -ORDER BY input.input_rank ASC, input.queued_at ASC, input.id ASC -FOR UPDATE; - --- name: GetNextQueuedAgentInputForAdmission :one -SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata -FROM agent_inputs input -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' -ORDER BY input.input_rank ASC, input.queued_at ASC, input.id ASC -FOR UPDATE -LIMIT 1; - -- name: ListQueuedBacklogInputs :many SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata FROM agent_inputs input @@ -141,309 +50,3 @@ WHERE input.project_id = sqlc.arg(project_id) ) ORDER BY input.delivery_mode DESC, input.input_rank ASC, input.queued_at ASC, input.id ASC LIMIT sqlc.arg(row_limit)::bigint; - --- name: AdmitAgentInput :one -UPDATE agent_inputs input -SET state = 'resolved', - admitted_event_id = sqlc.arg(admitted_event_id)::uuid, - admitted_at = event.created_at, - resolved_at = event.created_at -FROM agent_events event -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND event.agent_id = input.agent_id - AND event.agent_input_id = input.id - AND event.id = sqlc.arg(admitted_event_id) -RETURNING input.admitted_at, input.resolved_at; - --- name: CancelQueuedBacklogInput :execrows -UPDATE agent_inputs -SET state = 'canceled', - canceled_at = statement_timestamp() -WHERE agent_inputs.project_id = sqlc.arg(project_id) - AND agent_inputs.agent_id = sqlc.arg(agent_id) - AND agent_inputs.id = sqlc.arg(id) - AND agent_inputs.state = 'received' - AND agent_inputs.delivery_mode = 'queued' - AND agent_inputs.input_kind = 'content'; - --- name: CancelQueuedBacklogInputsForAgent :execrows -UPDATE agent_inputs -SET state = 'canceled', - canceled_at = statement_timestamp() -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content'; - --- name: QueuedBacklogMoveIsValid :one -SELECT CASE WHEN EXISTS ( - SELECT 1 - FROM agent_inputs input - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' -) AND ( - NOT sqlc.arg(requires_anchor)::boolean - OR EXISTS ( - SELECT 1 - FROM agent_inputs anchor - WHERE anchor.project_id = sqlc.arg(project_id) - AND anchor.agent_id = sqlc.arg(agent_id) - AND anchor.id = sqlc.arg(anchor_id) - AND anchor.state = 'received' - AND anchor.delivery_mode = 'queued' - AND anchor.input_kind = 'content' - ) -) THEN true ELSE false END AS valid; - --- name: MoveQueuedBacklogInputToFront :execrows -WITH first_neighbor AS ( - SELECT input_rank - FROM agent_inputs - WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' - AND id <> sqlc.arg(id) - ORDER BY input_rank ASC, queued_at ASC, id ASC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT input_rank FROM first_neighbor) IS NULL THEN sqlc.arg(rank_stride)::bigint - WHEN (SELECT input_rank FROM first_neighbor) > 1 THEN - (SELECT input_rank / 2 FROM first_neighbor) - ELSE NULL::bigint - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL; - --- name: MoveQueuedBacklogInputToBack :execrows -WITH last_neighbor AS ( - SELECT input_rank - FROM agent_inputs - WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' - AND id <> sqlc.arg(id) - ORDER BY input_rank DESC, queued_at DESC, id DESC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT input_rank FROM last_neighbor) IS NULL THEN sqlc.arg(rank_stride)::bigint - ELSE (SELECT input_rank + sqlc.arg(rank_stride)::bigint FROM last_neighbor) - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL; - --- name: MoveQueuedBacklogInputBefore :execrows -WITH target AS ( - SELECT anchor.id, anchor.input_rank, anchor.queued_at - FROM agent_inputs anchor - WHERE anchor.project_id = sqlc.arg(project_id) - AND anchor.agent_id = sqlc.arg(agent_id) - AND anchor.state = 'received' - AND anchor.delivery_mode = 'queued' - AND anchor.input_kind = 'content' - AND anchor.id = sqlc.arg(anchor_id) -), -previous_neighbor AS ( - SELECT candidate.input_rank - FROM agent_inputs candidate - CROSS JOIN target - WHERE candidate.project_id = sqlc.arg(project_id) - AND candidate.agent_id = sqlc.arg(agent_id) - AND candidate.state = 'received' - AND candidate.delivery_mode = 'queued' - AND candidate.input_kind = 'content' - AND candidate.id <> sqlc.arg(id) - AND (candidate.input_rank, candidate.queued_at, candidate.id) < (target.input_rank, target.queued_at, target.id) - ORDER BY candidate.input_rank DESC, candidate.queued_at DESC, candidate.id DESC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT id FROM target) IS NULL THEN NULL - WHEN (SELECT input_rank FROM previous_neighbor) IS NULL AND - (SELECT input_rank FROM target) > 1 THEN - (SELECT input_rank / 2 FROM target) - WHEN (SELECT input_rank FROM target) - (SELECT input_rank FROM previous_neighbor) > 1 THEN - (SELECT input_rank FROM previous_neighbor) + - ((SELECT input_rank FROM target) - (SELECT input_rank FROM previous_neighbor)) / 2 - ELSE NULL::bigint - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL; - --- name: MoveQueuedBacklogInputAfter :execrows -WITH target AS ( - SELECT anchor.id, anchor.input_rank, anchor.queued_at - FROM agent_inputs anchor - WHERE anchor.project_id = sqlc.arg(project_id) - AND anchor.agent_id = sqlc.arg(agent_id) - AND anchor.state = 'received' - AND anchor.delivery_mode = 'queued' - AND anchor.input_kind = 'content' - AND anchor.id = sqlc.arg(anchor_id) -), -next_neighbor AS ( - SELECT candidate.input_rank - FROM agent_inputs candidate - CROSS JOIN target - WHERE candidate.project_id = sqlc.arg(project_id) - AND candidate.agent_id = sqlc.arg(agent_id) - AND candidate.state = 'received' - AND candidate.delivery_mode = 'queued' - AND candidate.input_kind = 'content' - AND candidate.id <> sqlc.arg(id) - AND (candidate.input_rank, candidate.queued_at, candidate.id) > (target.input_rank, target.queued_at, target.id) - ORDER BY candidate.input_rank ASC, candidate.queued_at ASC, candidate.id ASC - LIMIT 1 -), -new_rank AS ( - SELECT CASE - WHEN (SELECT id FROM target) IS NULL THEN NULL - WHEN (SELECT input_rank FROM next_neighbor) IS NULL THEN - (SELECT input_rank + sqlc.arg(rank_stride)::bigint FROM target) - WHEN (SELECT input_rank FROM next_neighbor) - (SELECT input_rank FROM target) > 1 THEN - (SELECT input_rank FROM target) + - ((SELECT input_rank FROM next_neighbor) - (SELECT input_rank FROM target)) / 2 - ELSE NULL::bigint - END AS input_rank -) -UPDATE agent_inputs input -SET input_rank = new_rank.input_rank -FROM new_rank -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - AND new_rank.input_rank IS NOT NULL; - --- name: PromoteQueuedInputToSteering :one -WITH promoted AS ( - UPDATE agent_inputs input - SET delivery_mode = 'steering', - input_rank = coalesce( - ( - SELECT max(existing.input_rank) + sqlc.arg(rank_stride)::bigint - FROM agent_inputs existing - WHERE existing.project_id = input.project_id - AND existing.agent_id = input.agent_id - AND existing.delivery_mode = 'steering' - AND existing.state = 'received' - AND existing.input_kind = 'content' - ), - sqlc.arg(rank_stride)::bigint - ) - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'queued' - AND input.input_kind = 'content' - RETURNING TRUE -) -SELECT EXISTS (SELECT 1 FROM promoted) AS changed, - EXISTS (SELECT 1 FROM promoted) - OR EXISTS ( - SELECT 1 - FROM agent_inputs input - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.input_kind = 'content' - AND ( - (input.state = 'received' AND input.delivery_mode = 'steering') - OR (input.state = 'resolved' AND input.admitted_event_id IS NOT NULL) - ) - ) AS effective; - --- name: DemoteSteeringInputToQueued :execrows -UPDATE agent_inputs input -SET delivery_mode = 'queued', - input_rank = coalesce( - ( - SELECT max(existing.input_rank) + sqlc.arg(rank_stride)::bigint - FROM agent_inputs existing - WHERE existing.project_id = input.project_id - AND existing.agent_id = input.agent_id - AND existing.delivery_mode = 'queued' - AND existing.state = 'received' - AND existing.input_kind = 'content' - ), - sqlc.arg(rank_stride)::bigint - ) -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - AND input.delivery_mode = 'steering' - AND input.input_kind = 'content'; - --- name: RebalanceQueuedBacklogRanks :execrows -WITH ordered AS ( - SELECT id, row_number() OVER (ORDER BY input_rank ASC, queued_at ASC, id ASC) AS row_number - FROM agent_inputs - WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND state = 'received' - AND delivery_mode = 'queued' - AND input_kind = 'content' -) -UPDATE agent_inputs input -SET input_rank = ordered.row_number * sqlc.arg(rank_stride)::bigint -FROM ordered -WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = ordered.id; - --- name: CancelSteeringAgentInputsForAgent :execrows -UPDATE agent_inputs -SET state = 'canceled', - canceled_at = statement_timestamp() -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND state = 'received' - AND delivery_mode = 'steering' - AND input_kind = 'content'; diff --git a/internal/storage/queries/agent_inputs_ingress.sql b/internal/storage/queries/agent_inputs_ingress.sql index d5c9b83db..20e12e743 100644 --- a/internal/storage/queries/agent_inputs_ingress.sql +++ b/internal/storage/queries/agent_inputs_ingress.sql @@ -1,31 +1,3 @@ --- name: InsertAgentInput :one -WITH generated AS ( - SELECT coalesce(sqlc.narg(id), uuidv7()) AS id -) -INSERT INTO agent_inputs(id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, delivery_mode, idempotency_scope, input_idempotency_key, queued_at, metadata) -SELECT generated.id, agent.project_id, agent.id, 'received', - coalesce( - ( - SELECT max(existing.input_rank) + sqlc.arg(rank_stride)::bigint - FROM agent_inputs existing - WHERE existing.project_id = agent.project_id - AND existing.agent_id = agent.id - AND existing.delivery_mode = sqlc.arg(delivery_mode)::text - AND existing.state = 'received' - AND existing.input_kind = 'content' - ), - sqlc.arg(rank_stride)::bigint - ), - sqlc.arg(actor_id), 'content', - sqlc.narg(integration_target_id)::uuid, - sqlc.arg(delivery_mode)::text, - sqlc.narg(idempotency_scope), sqlc.narg(input_idempotency_key), statement_timestamp(), sqlc.arg(metadata) -FROM agents agent -JOIN generated ON true -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) -RETURNING id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata; - -- name: GetAgentInputByIdempotency :one SELECT id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata FROM agent_inputs @@ -33,3 +5,10 @@ WHERE project_id = sqlc.arg(project_id) AND agent_id = sqlc.arg(agent_id) AND idempotency_scope = sqlc.arg(idempotency_scope)::text AND input_idempotency_key = sqlc.arg(input_idempotency_key)::text; + +-- name: GetAgentInput :one +SELECT id, project_id, agent_id, state, input_rank, actor_id, input_kind, integration_target_id, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata +FROM agent_inputs +WHERE project_id = sqlc.arg(project_id) + AND agent_id = sqlc.arg(agent_id) + AND id = sqlc.arg(id); diff --git a/internal/storage/queries/agent_interactions.sql b/internal/storage/queries/agent_interactions.sql index 6f3643789..1d31b2aa4 100644 --- a/internal/storage/queries/agent_interactions.sql +++ b/internal/storage/queries/agent_interactions.sql @@ -1,17 +1,3 @@ --- name: InsertAgentInteraction :one -INSERT INTO agent_interactions( - agent_id, tool_call_id, - interaction_kind, state, request, created_at, destination -) -SELECT tool_call.agent_id, - tool_call.id, sqlc.arg(interaction_kind), 'open', - sqlc.arg(request), statement_timestamp(), sqlc.narg(destination)::jsonb -FROM tool_call_read_projection tool_call -WHERE tool_call.project_id = sqlc.arg(project_id) - AND tool_call.agent_id = sqlc.arg(agent_id) - AND tool_call.id = sqlc.arg(tool_call_id) -RETURNING id; - -- name: GetAgentInteraction :one SELECT id, project_id, agent_id, turn_id, model_call_context_id, tool_call_id, provider_call_id, @@ -22,31 +8,6 @@ WHERE project_id = sqlc.arg(project_id) AND agent_id = sqlc.arg(agent_id) AND id = sqlc.arg(id); --- name: ListAgentInteractionsByIDs :many -SELECT id, project_id, agent_id, turn_id, - model_call_context_id, tool_call_id, provider_call_id, - interaction_kind, state, request, resolution, - resolved_by_input_id, created_at, resolved_at, destination, presentation_receipt -FROM agent_interaction_read_projection -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND id = ANY(sqlc.arg(ids)::uuid[]) -ORDER BY created_at, id; - --- name: AgentInteractionHasLaterStop :one -SELECT EXISTS ( - SELECT 1 -FROM agent_interactions interaction - JOIN tool_call_read_projection tool_call ON tool_call.agent_id = interaction.agent_id - AND tool_call.id = interaction.tool_call_id - JOIN agent_stop_events stop_event ON stop_event.project_id = tool_call.project_id - AND stop_event.agent_id = tool_call.agent_id - AND stop_event.sequence > tool_call.source_event_sequence - WHERE tool_call.project_id = sqlc.arg(project_id) - AND interaction.agent_id = sqlc.arg(agent_id) - AND interaction.id = sqlc.arg(id) -) AS stopped; - -- name: GetAgentInteractionByToolCallKind :one SELECT id, project_id, agent_id, turn_id, model_call_context_id, tool_call_id, provider_call_id, @@ -57,91 +18,3 @@ WHERE project_id = sqlc.arg(project_id) AND agent_id = sqlc.arg(agent_id) AND tool_call_id = sqlc.arg(tool_call_id) AND interaction_kind = sqlc.arg(interaction_kind); - --- name: ResolveAgentInteraction :one -UPDATE agent_interactions interaction -SET state = 'resolved', - resolution = sqlc.arg(resolution), - resolved_by_input_id = sqlc.narg(resolved_by_input_id)::uuid, - resolved_at = statement_timestamp() -WHERE interaction.agent_id = sqlc.arg(agent_id) - AND interaction.id = sqlc.arg(id) - AND interaction.state = 'open' - AND EXISTS ( - SELECT 1 - FROM agents agent - WHERE agent.id = interaction.agent_id - AND agent.project_id = sqlc.arg(project_id) - ) -RETURNING interaction.id; - --- name: InsertInteractionResponseAgentInput :one -WITH target AS ( - SELECT projection.project_id, interaction.agent_id, interaction.id - FROM agent_interactions interaction - JOIN agent_interaction_read_projection projection - ON projection.agent_id = interaction.agent_id - AND projection.id = interaction.id - WHERE projection.project_id = sqlc.arg(project_id) - AND interaction.agent_id = sqlc.arg(agent_id) - AND interaction.id = sqlc.arg(target_interaction_id) -) -INSERT INTO agent_inputs( - project_id, agent_id, state, - actor_id, input_kind, delivery_mode, target_interaction_id, - idempotency_scope, input_idempotency_key, queued_at, metadata -) -SELECT target.project_id, target.agent_id, 'received', - sqlc.narg(actor_id)::uuid, - 'interaction_response', 'immediate', target.id, - sqlc.arg(idempotency_scope), sqlc.arg(input_idempotency_key), - statement_timestamp(), sqlc.arg(metadata) -FROM target -ON CONFLICT (project_id, agent_id, idempotency_scope, input_idempotency_key) WHERE idempotency_scope IS NOT NULL AND input_idempotency_key IS NOT NULL DO NOTHING -RETURNING id, project_id, agent_id, state, input_rank, actor_id, input_kind, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata; - --- name: ResolveInteractionResponseAgentInput :one -UPDATE agent_inputs -SET state = 'resolved', - admitted_event_id = sqlc.arg(event_id), - admitted_at = statement_timestamp(), - resolved_at = statement_timestamp() -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND id = sqlc.arg(id) - AND input_kind = 'interaction_response' - AND target_interaction_id = sqlc.arg(target_interaction_id) - AND state = 'received' -RETURNING id, project_id, agent_id, state, input_rank, actor_id, input_kind, coalesce(idempotency_scope, '') AS idempotency_scope, coalesce(input_idempotency_key, '') AS input_idempotency_key, queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, coalesce(rejected_reason, '') AS rejected_reason, metadata; - --- name: CancelOpenAgentInteractionsForAgent :many -UPDATE agent_interactions interaction -SET state = 'canceled', - resolution = jsonb_build_object('reason', sqlc.arg(reason)::text), - resolved_by_input_id = sqlc.narg(resolved_by_input_id)::uuid, - resolved_at = statement_timestamp() -FROM agent_interaction_read_projection projection -WHERE interaction.agent_id = sqlc.arg(agent_id) - AND interaction.state = 'open' - AND projection.project_id = sqlc.arg(project_id) - AND projection.agent_id = interaction.agent_id - AND projection.id = interaction.id - AND projection.turn_id = sqlc.arg(turn_id) -RETURNING interaction.id; - --- name: CancelOpenAgentInteractionsForToolCall :many -UPDATE agent_interactions -SET state = 'canceled', - resolution = jsonb_build_object('reason', sqlc.arg(reason)::text), - resolved_by_input_id = NULL, - resolved_at = statement_timestamp() -WHERE agent_interactions.agent_id = sqlc.arg(agent_id) - AND agent_interactions.tool_call_id = sqlc.arg(tool_call_id) - AND agent_interactions.state = 'open' - AND EXISTS ( - SELECT 1 - FROM agents agent - WHERE agent.id = agent_interactions.agent_id - AND agent.project_id = sqlc.arg(project_id) - ) -RETURNING id, interaction_kind, state; diff --git a/internal/storage/queries/agent_machine_bindings.sql b/internal/storage/queries/agent_machine_bindings.sql index 6566a45da..a8384a464 100644 --- a/internal/storage/queries/agent_machine_bindings.sql +++ b/internal/storage/queries/agent_machine_bindings.sql @@ -178,21 +178,6 @@ SELECT pg_advisory_xact_lock( hashtextextended('agent_machine_sources:' || sqlc.arg(agent_id)::uuid::text, 0) ); --- name: LockAttachedAgentPoolMachines :exec -SELECT machine.id -FROM machines machine -JOIN agent_machine_bindings binding ON binding.org_id = machine.org_id - AND binding.machine_id = machine.id -WHERE binding.project_id = sqlc.arg(project_id) - AND binding.agent_id = sqlc.arg(agent_id) - AND binding.binding_kind = 'pool' - AND binding.state = 'attached' - AND machine.source_kind = 'pool' - AND machine.deleted_at IS NULL - AND machine.lifecycle_state NOT IN ('deleting', 'delete_failed', 'deleted') -ORDER BY machine.id -FOR UPDATE OF machine; - -- name: ListAttachedAgentPoolMachineIDsForLifecycle :many SELECT machine.id FROM machines machine diff --git a/internal/storage/queries/agent_runtime.sql b/internal/storage/queries/agent_runtime.sql index 34fcaa74a..4132b5b45 100644 --- a/internal/storage/queries/agent_runtime.sql +++ b/internal/storage/queries/agent_runtime.sql @@ -1,48 +1,3 @@ --- name: InsertAgent :one --- @sqlc-vet-disable configured-models-deleted-at model-provider-configs-deleted-at --- Display-only model names must still resolve after the model or provider config is soft deleted. -WITH inserted AS ( - INSERT INTO agents( - id, org_id, project_id, state, name, agent_profile_id, current_config_id, - idempotency_key, parent_agent_id, subagent_key, - archive_after_idle_minutes, created_at, updated_at - ) - SELECT - coalesce(sqlc.narg(id)::uuid, uuidv7()), sqlc.arg(org_id), sqlc.arg(project_id), 'active', sqlc.arg(name), - sqlc.narg(agent_profile_id), sqlc.arg(current_config_id), sqlc.narg(idempotency_key), - sqlc.narg(parent_agent_id), sqlc.arg(subagent_key), - sqlc.narg(archive_after_idle_minutes), - transaction_timestamp(), transaction_timestamp() - FROM projects project - JOIN orgs org ON org.id = project.org_id - WHERE project.org_id = sqlc.arg(org_id) - AND project.id = sqlc.arg(project_id) - AND project.deleted_at IS NULL - AND org.deleted_at IS NULL - ON CONFLICT (project_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING - RETURNING id, org_id, project_id, state, name, - agent_profile_id, current_config_id, interaction_target_id, - idempotency_key, next_event_sequence, created_at, updated_at, archived_at, - parent_agent_id, subagent_key -) -SELECT agent.id, agent.org_id, agent.project_id, agent.state, agent.name, - agent.agent_profile_id, agent.current_config_id, agent.interaction_target_id, - coalesce(agent.idempotency_key, '') AS idempotency_key, - agent.next_event_sequence, agent.created_at, agent.updated_at, agent.archived_at, - agent.parent_agent_id, agent.subagent_key, - coalesce(configured_model.name, '') AS model_name, - coalesce(model_provider_config.name, '') AS model_provider_config_name -FROM inserted agent -LEFT JOIN agent_configs agent_config - ON agent_config.project_id = agent.project_id - AND agent_config.id = agent.current_config_id -LEFT JOIN configured_models configured_model - ON configured_model.org_id = agent.org_id - AND configured_model.id = agent_config.configured_model_id -LEFT JOIN model_provider_configs model_provider_config - ON model_provider_config.org_id = configured_model.org_id - AND model_provider_config.id = configured_model.model_provider_config_id; - -- name: LockAgentLaunchIdempotencyKey :exec SELECT pg_advisory_xact_lock( hashtextextended( @@ -282,19 +237,6 @@ WHERE agent.project_id = ANY(sqlc.arg(project_ids)::uuid[]) ORDER BY agent.updated_at DESC, agent.id DESC LIMIT sqlc.arg(row_limit)::bigint; --- name: ArchiveAgent :execrows --- Idempotent: archiving an already-archived agent matches the row and keeps --- the original archived_at. -UPDATE agents -SET state = 'archived', - archived_at = coalesce(archived_at, statement_timestamp()), - updated_at = CASE - WHEN state = 'archived' THEN updated_at - ELSE statement_timestamp() - END -WHERE project_id = sqlc.arg(project_id) - AND id = sqlc.arg(id); - -- name: AgentExistsInProject :one SELECT EXISTS ( SELECT 1 diff --git a/internal/storage/queries/agent_runtime_locks.sql b/internal/storage/queries/agent_runtime_locks.sql index b395aa9cd..f36bbb242 100644 --- a/internal/storage/queries/agent_runtime_locks.sql +++ b/internal/storage/queries/agent_runtime_locks.sql @@ -1,74 +1,6 @@ -- name: DBNow :one SELECT transaction_timestamp()::timestamptz; --- name: AcquireAgentRuntimeLock :one -WITH lease_clock AS MATERIALIZED ( - SELECT statement_timestamp() AS renewed_at -) -INSERT INTO agent_runtime_locks(agent_id, worker_process_id, started_at, renewed_at, lease_expires_at) -SELECT agent.id, - sqlc.arg(worker_process_id)::uuid, - lease_clock.renewed_at, - lease_clock.renewed_at, - lease_clock.renewed_at + (sqlc.arg(lease_duration_microseconds)::bigint * interval '1 microsecond') -FROM agents agent -CROSS JOIN lease_clock -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - AND agent.state <> 'archived' -ON CONFLICT (agent_id) DO NOTHING -RETURNING id, agent_id, worker_process_id, started_at, renewed_at, - lease_expires_at, cancel_requested_at; - --- name: LockAgentRuntimeLockForRenewal :one -SELECT runtime_lock.id -FROM agent_runtime_locks runtime_lock -JOIN agents agent ON agent.id = runtime_lock.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(id) -FOR UPDATE OF runtime_lock; - --- name: RenewAgentRuntimeLock :one -UPDATE agent_runtime_locks runtime_lock -SET renewed_at = greatest(runtime_lock.renewed_at, lease_clock.statement_at), - lease_expires_at = greatest( - runtime_lock.lease_expires_at, - greatest(runtime_lock.renewed_at, lease_clock.statement_at) + - (sqlc.arg(lease_duration_microseconds)::bigint * interval '1 microsecond') - ) -FROM (SELECT statement_timestamp() AS statement_at) AS lease_clock -CROSS JOIN agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = runtime_lock.agent_id - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(id) -RETURNING runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, - runtime_lock.started_at, runtime_lock.renewed_at, - runtime_lock.lease_expires_at, runtime_lock.cancel_requested_at; - --- name: RequestAgentRuntimeCancel :one -UPDATE agent_runtime_locks runtime_lock -SET cancel_requested_at = coalesce(runtime_lock.cancel_requested_at, greatest(runtime_lock.started_at, statement_timestamp())) -FROM agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = runtime_lock.agent_id - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.cancel_requested_at IS NULL -RETURNING runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, - runtime_lock.started_at, runtime_lock.renewed_at, - runtime_lock.lease_expires_at, runtime_lock.cancel_requested_at; - --- name: GetPendingAgentRuntimeCancel :one -SELECT runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, - runtime_lock.started_at, runtime_lock.renewed_at, - runtime_lock.lease_expires_at, runtime_lock.cancel_requested_at -FROM agent_runtime_locks runtime_lock -JOIN agents agent ON agent.id = runtime_lock.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.cancel_requested_at IS NOT NULL; - -- name: LockAgentRuntimeLockForOwnedMutation :one SELECT runtime_lock.id FROM agent_runtime_locks runtime_lock @@ -78,19 +10,6 @@ WHERE agent.project_id = sqlc.arg(project_id) AND runtime_lock.id = sqlc.arg(id) FOR UPDATE OF runtime_lock; --- name: AgentRuntimeLockIsActive :one -SELECT EXISTS ( - SELECT 1 - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - AND agent.state <> 'archived' - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) AS active; - -- name: GetAgentRuntimeLockForRelease :one SELECT runtime_lock.id, runtime_lock.agent_id, runtime_lock.worker_process_id, runtime_lock.started_at, runtime_lock.renewed_at, @@ -101,16 +20,6 @@ WHERE agent.project_id = sqlc.arg(project_id) AND runtime_lock.agent_id = sqlc.arg(agent_id) AND runtime_lock.id = sqlc.arg(id); --- name: ReleaseAgentRuntimeLock :execrows -DELETE FROM agent_runtime_locks -USING agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = agent_runtime_locks.agent_id - AND agent_runtime_locks.agent_id = sqlc.arg(agent_id) - AND agent_runtime_locks.id = sqlc.arg(id); - --- Candidate discovery intentionally does not lock runtime rows. Destructive --- recovery must lock the canonical agent row before the runtime row. -- name: ListExpiredAgentRuntimeLockCandidates :many SELECT runtime_lock.id, agent.project_id, runtime_lock.agent_id FROM agent_runtime_locks runtime_lock @@ -119,24 +28,6 @@ WHERE runtime_lock.lease_expires_at <= statement_timestamp() ORDER BY runtime_lock.lease_expires_at ASC, runtime_lock.id ASC LIMIT sqlc.arg(batch_size)::int; --- name: LockExpiredAgentRuntimeLockForReap :one -SELECT runtime_lock.id, agent.project_id, runtime_lock.agent_id -FROM agent_runtime_locks runtime_lock -JOIN agents agent ON agent.id = runtime_lock.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(id) - AND runtime_lock.lease_expires_at <= statement_timestamp() -FOR UPDATE OF runtime_lock SKIP LOCKED; - --- name: DeleteAgentRuntimeLockForReap :execrows -DELETE FROM agent_runtime_locks -USING agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = agent_runtime_locks.agent_id - AND agent_runtime_locks.agent_id = sqlc.arg(agent_id) - AND agent_runtime_locks.id = sqlc.arg(id); - -- name: FailQueuedProcessesForRuntimeEnd :many UPDATE processes process SET state = 'failed', @@ -169,24 +60,3 @@ WHERE action.project_id = sqlc.arg(project_id) AND tool_call.state = 'waiting' ) RETURNING action.id, action.org_id, action.project_id, action.agent_id, action.process_id, action.tool_call_id, action.runtime_lock_id, action.action_kind, action.seq, action.payload, action.state, action.created_at, action.updated_at, action.state_reason_code, action.state_reason_message; - --- name: FailRuntimeToolCalls :many -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM tool_call_read_projection projection -WHERE call.agent_id = sqlc.arg(agent_id) - AND call.runtime_lock_id = sqlc.arg(runtime_lock_id)::uuid - AND call.state = 'running' - AND call.type IN ('built_in', 'mcp') - AND projection.project_id = sqlc.arg(project_id) - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, - projection.source_event_id, projection.model_call_context_id, call.provider_call_id, - call.name, call.input, - call.type, call.state, - 'failed'::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; diff --git a/internal/storage/queries/agent_turns.sql b/internal/storage/queries/agent_turns.sql index c0d080a1f..1e1abb8b9 100644 --- a/internal/storage/queries/agent_turns.sql +++ b/internal/storage/queries/agent_turns.sql @@ -1,10 +1,3 @@ --- name: NextTurnSequence :one -SELECT (coalesce(max(turn.turn_sequence), 0) + 1)::bigint -FROM agent_turns turn -JOIN agents agent ON agent.id = turn.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND turn.agent_id = sqlc.arg(agent_id); - -- name: ListAgentTurnsForRead :many SELECT turn.id, agent.project_id, @@ -141,68 +134,6 @@ WHERE projection.project_id = sqlc.arg(project_id) ) ORDER BY projection.turn_sequence DESC, projection.sequence ASC; --- name: InsertAgentTurn :one -WITH inserted AS ( - INSERT INTO agent_turns(id, agent_id, turn_sequence, latest_event_id, latest_semantic_event_id) - SELECT sqlc.arg(id), sqlc.arg(agent_id), - sqlc.arg(turn_sequence), event.id, sqlc.arg(latest_semantic_event_id) - FROM agent_events event - JOIN agents agent ON agent.id = event.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND event.agent_id = sqlc.arg(agent_id) - AND event.turn_id = sqlc.arg(id) - AND event.id = sqlc.arg(latest_event_id) - RETURNING id, agent_id, turn_sequence, latest_event_id, latest_semantic_event_id -) -SELECT inserted.id, agent.project_id, inserted.agent_id, inserted.turn_sequence, - inserted.latest_event_id, inserted.latest_semantic_event_id -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id; - --- name: UpdateAgentTurnLatestEvent :execrows -UPDATE agent_turns -SET latest_event_id = sqlc.arg(latest_event_id), - latest_semantic_event_id = coalesce(sqlc.narg(latest_semantic_event_id), latest_semantic_event_id) -FROM agent_events event -WHERE agent_turns.agent_id = sqlc.arg(agent_id) - AND agent_turns.id = sqlc.arg(id) - AND event.agent_id = agent_turns.agent_id - AND event.turn_id = agent_turns.id - AND event.id = sqlc.arg(latest_event_id) - AND EXISTS ( - SELECT 1 - FROM agents agent - WHERE agent.id = agent_turns.agent_id - AND agent.project_id = sqlc.arg(project_id) - ); - --- name: CurrentContinuableAgentTurn :one -WITH latest_turn AS MATERIALIZED ( - SELECT turn.id, agent.project_id, turn.agent_id, turn.turn_sequence, turn.latest_event_id, - turn.latest_semantic_event_id - FROM agent_turns turn - JOIN agents agent ON agent.id = turn.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND turn.agent_id = sqlc.arg(agent_id) - ORDER BY turn.turn_sequence DESC - LIMIT 1 -) -SELECT turn.id, turn.project_id, turn.agent_id, turn.turn_sequence, turn.latest_event_id, - turn.latest_semantic_event_id -FROM latest_turn turn -WHERE EXISTS ( - SELECT 1 - FROM agent_continuable_model_contexts(turn.project_id, turn.agent_id) - AS context(turn_id, model_call_context_id, input_event_sequence, has_later_semantic_event) - WHERE context.turn_id = turn.id - ) - OR agent_has_incomplete_tool_batch(turn.project_id, turn.agent_id) - OR EXISTS ( - SELECT 1 - FROM agent_next_model_work(turn.project_id, turn.agent_id) frontier - WHERE frontier.turn_id = turn.id - ); - -- name: IsOutputLimitBoundary :one SELECT EXISTS ( SELECT 1 @@ -294,3 +225,8 @@ GROUP BY event.id, event.sequence, event.created_at, event.event_kind, output.provider_replay, output.stop_reason, input.input_kind ORDER BY event.sequence ASC LIMIT sqlc.arg(page_limit); + +-- name: GetAgentTurn :one +SELECT t.id,t.turn_sequence,t.latest_event_id,t.latest_semantic_event_id +FROM agent_turns t JOIN agents a ON a.id=t.agent_id +WHERE a.project_id=sqlc.arg(project_id) AND t.agent_id=sqlc.arg(agent_id) AND t.id=sqlc.arg(id); diff --git a/internal/storage/queries/agent_work.sql b/internal/storage/queries/agent_work.sql deleted file mode 100644 index aca76bb72..000000000 --- a/internal/storage/queries/agent_work.sql +++ /dev/null @@ -1,49 +0,0 @@ --- name: NextAgentToolWork :one -SELECT frontier.turn_id::uuid AS turn_id, - frontier.model_call_context_id::uuid AS model_call_context_id, - frontier.model_output_id::uuid AS model_output_id, - frontier.source_event_id::uuid AS source_event_id -FROM agent_tool_work_frontiers( - sqlc.arg(project_id), - sqlc.arg(agent_id) -) AS frontier( - turn_id, - model_call_context_id, - model_output_id, - source_event_id, - ready_at, - frontier_order_created_at, - frontier_order_tool_call_id -) -ORDER BY frontier.ready_at, - frontier.frontier_order_created_at, - frontier.frontier_order_tool_call_id -LIMIT 1; - --- name: NextAgentModelWork :one -SELECT frontier.work_kind::text AS work_kind, - coalesce(frontier.model_call_context_id, '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS model_call_context_id, - coalesce(frontier.model_output_id, '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS model_output_id, - frontier.turn_id::uuid AS turn_id, - frontier.input_ids::uuid[] AS input_ids, - frontier.opening_event_sequence::bigint AS opening_event_sequence, - frontier.ready_at <= statement_timestamp() AS is_ready -FROM agent_next_model_work( - sqlc.arg(project_id), - sqlc.arg(agent_id) -) AS frontier( - work_kind, - model_call_context_id, - model_output_id, - turn_id, - input_ids, - opening_event_sequence, - ready_at -) -LIMIT 1; - --- name: AgentHasIncompleteToolBatch :one -SELECT agent_has_incomplete_tool_batch( - sqlc.arg(project_id), - sqlc.arg(agent_id) -) AS incomplete; diff --git a/internal/storage/queries/agents.sql b/internal/storage/queries/agents.sql index 54aa808f6..632756f07 100644 --- a/internal/storage/queries/agents.sql +++ b/internal/storage/queries/agents.sql @@ -170,13 +170,6 @@ JOIN inserted_version version AND version.profile_id = profile.id AND version.id = profile.current_version_id; --- name: UpdateAgentCurrentConfig :execrows -UPDATE agents -SET current_config_id = sqlc.arg(agent_config_id), - updated_at = statement_timestamp() -WHERE project_id = sqlc.arg(project_id) - AND id = sqlc.arg(agent_id); - -- name: GetAgentProfile :one SELECT profile.id, project.org_id, profile.project_id, profile.name, version.agent_config_id AS current_config_id, diff --git a/internal/storage/queries/context_checkpoints.sql b/internal/storage/queries/context_checkpoints.sql index 519c03d0b..a337596e9 100644 --- a/internal/storage/queries/context_checkpoints.sql +++ b/internal/storage/queries/context_checkpoints.sql @@ -1,51 +1,3 @@ --- name: CountCheckpointRangeEvents :one -SELECT count(*)::bigint AS count -FROM agent_events -JOIN agents agent ON agent.id = agent_events.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND agent_events.agent_id = sqlc.arg(agent_id) - AND sequence BETWEEN sqlc.arg(start_sequence) AND sqlc.arg(end_sequence); - --- name: CountCheckpointRangeOpenToolResults :one -SELECT count(*)::bigint AS count -FROM tool_call_read_projection tool_call -LEFT JOIN tool_call_results result ON result.agent_id = tool_call.agent_id - AND result.tool_call_id = tool_call.id -LEFT JOIN agent_events result_event ON result_event.agent_id = result.agent_id - AND result_event.tool_call_result_id = result.id - AND result_event.event_kind = 'tool_result' -WHERE tool_call.project_id = sqlc.arg(project_id) - AND tool_call.agent_id = sqlc.arg(agent_id) - AND tool_call.source_event_sequence BETWEEN sqlc.arg(start_sequence) AND sqlc.arg(end_sequence) - AND ( - result_event.id IS NULL - OR result_event.sequence > sqlc.arg(end_sequence) - ); - --- name: InsertContextCheckpoint :one -INSERT INTO context_checkpoints( - id, agent_id, summarized_through_event_sequence, - producer_model_call_context_id, summary, created_at -) -SELECT sqlc.arg(id), agent.id, - context.source_event_sequence_end, - context.id, - sqlc.arg(summary), - statement_timestamp() -FROM agents agent -JOIN model_call_contexts context ON context.project_id = agent.project_id - AND context.agent_id = agent.id - AND context.id = sqlc.arg(producer_model_call_context_id) - AND context.operation_kind = 'compaction' - AND context.state = 'started' -JOIN agent_runtime_locks runtime_lock ON runtime_lock.agent_id = context.agent_id - AND runtime_lock.id = context.runtime_lock_id - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) -RETURNING id; - -- name: GetContextCheckpoint :one SELECT checkpoint.id, agent.project_id, checkpoint.agent_id, checkpoint.summarized_through_event_sequence, diff --git a/internal/storage/queries/events.sql b/internal/storage/queries/events.sql index a0e2ac6d2..496f5c347 100644 --- a/internal/storage/queries/events.sql +++ b/internal/storage/queries/events.sql @@ -4,26 +4,6 @@ FROM agent_events event JOIN agents agent ON agent.id = event.agent_id WHERE agent.project_id = sqlc.arg(project_id) AND event.agent_id = sqlc.arg(agent_id) AND event.idempotency_key = sqlc.arg(idempotency_key)::text; --- name: AllocateEventSequence :one -SELECT project_id, next_event_sequence -FROM agents AS agents -WHERE id = $1 -FOR UPDATE; - --- name: LatestAgentEvent :one -SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, event.sequence, event.event_kind, event.created_at, coalesce(event.idempotency_key, '') AS idempotency_key -FROM agent_events event -JOIN agents agent ON agent.id = event.agent_id -WHERE agent.project_id = $1 AND event.agent_id = $2 -ORDER BY sequence DESC -LIMIT 1; - --- name: AdvanceEventSequence :exec -UPDATE agents -SET next_event_sequence = next_event_sequence + 1, updated_at = statement_timestamp() -WHERE project_id = sqlc.arg(project_id) - AND id = sqlc.arg(id); - -- name: MaxEventSequence :one SELECT coalesce(max(event.sequence), 0)::bigint FROM agent_events event diff --git a/internal/storage/queries/identity.sql b/internal/storage/queries/identity.sql index 01de1e061..1994b086c 100644 --- a/internal/storage/queries/identity.sql +++ b/internal/storage/queries/identity.sql @@ -1251,11 +1251,6 @@ FROM orgs WHERE id = $1 AND deleted_at IS NULL FOR UPDATE; --- name: CountOrgMemberships :one -SELECT count(*)::bigint -FROM org_memberships -WHERE org_id = $1; - -- name: ListProjectAuthorizationRolesForPrincipal :many SELECT role FROM principal_project_authorization_roles @@ -1496,3 +1491,8 @@ SELECT encrypted_private_key FROM oidc_signing_keys WHERE id = 'default'; INSERT INTO oidc_signing_keys(id, encrypted_private_key) VALUES ('default', sqlc.arg(encrypted_private_key)) ON CONFLICT (id) DO NOTHING; + +-- name: CountOrgMemberships :one +SELECT count(*)::bigint +FROM org_memberships +WHERE org_id = $1; diff --git a/internal/storage/queries/integration_interactions.sql b/internal/storage/queries/integration_interactions.sql index b1acdab8c..720999ee5 100644 --- a/internal/storage/queries/integration_interactions.sql +++ b/internal/storage/queries/integration_interactions.sql @@ -10,24 +10,6 @@ FROM agent_interaction_read_projection WHERE project_id = sqlc.arg(project_id) AND id = sqlc.arg(interaction_id) AND destination ->> 'integration_id' = sqlc.arg(integration_id)::text; --- name: SetInteractionSelection :execrows -UPDATE agents -SET interaction_target_id = sqlc.narg(target_id)::uuid, - interaction_handler_key = sqlc.narg(handler_key)::text, - interaction_auto_select = sqlc.arg(auto_select)::boolean, - updated_at = statement_timestamp() -WHERE agents.project_id = sqlc.arg(project_id) AND agents.id = sqlc.arg(agent_id) - AND ((sqlc.narg(target_id)::uuid IS NULL AND sqlc.narg(handler_key)::text IS NULL) - OR (sqlc.narg(handler_key)::text <> '' AND EXISTS ( - SELECT 1 FROM integration_targets target - JOIN integrations integration - ON integration.project_id = target.project_id - AND integration.id = target.integration_id - WHERE target.project_id = agents.project_id AND target.agent_id = agents.id - AND target.id = sqlc.narg(target_id)::uuid AND target.deleted_at IS NULL - AND integration.deleted_at IS NULL AND integration.state = 'active' - ))); - -- name: GetInteractionDestinationTarget :one SELECT target.id, target.integration_id AS integration_id, target.scope_kind, target.scope_ref, target.display_name, @@ -39,17 +21,6 @@ WHERE target.project_id = sqlc.arg(project_id) AND target.agent_id = sqlc.arg(ag AND target.id = sqlc.arg(target_id) AND target.deleted_at IS NULL AND integration.deleted_at IS NULL; --- name: RecordAgentInteractionPresentationReceipt :execrows -UPDATE agent_interactions interaction -SET presentation_receipt = sqlc.arg(receipt)::jsonb -WHERE interaction.agent_id = sqlc.arg(agent_id) AND interaction.id = sqlc.arg(id) - AND interaction.destination = sqlc.arg(destination)::jsonb - AND interaction.presentation_receipt IS NULL - AND EXISTS ( - SELECT 1 FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) AND agent.id = interaction.agent_id - ); - -- name: GetInteractionCallbackIntegrationID :one SELECT (destination ->> 'integration_id')::uuid AS integration_id FROM agent_interactions diff --git a/internal/storage/queries/integration_target_lifecycle.sql b/internal/storage/queries/integration_target_lifecycle.sql index 023770c41..5dea16970 100644 --- a/internal/storage/queries/integration_target_lifecycle.sql +++ b/internal/storage/queries/integration_target_lifecycle.sql @@ -12,16 +12,6 @@ SELECT subscription.agent_id FROM integration_subscriptions subscription WHERE subscription.project_id = sqlc.arg(project_id) AND subscription.integration_id = sqlc.arg(integration_id) ORDER BY agent_id; --- name: ClearDeletedIntegrationTargetsFromAgents :exec --- @sqlc-vet-disable integration-targets-deleted-at -UPDATE agents agent SET interaction_target_id = NULL, interaction_handler_key = NULL, updated_at = statement_timestamp() -FROM integration_targets target -WHERE agent.project_id = sqlc.arg(project_id) - AND target.project_id = agent.project_id - AND target.integration_id = sqlc.arg(integration_id) - AND agent.id = target.agent_id - AND agent.interaction_target_id = target.id; - -- name: GetIntegrationTarget :one SELECT target.id, project.org_id, target.project_id, target.agent_id, target.integration_id, target.scope_ref, target.scope_kind, target.display_name, target.launch_key, target.deleted_at, target.created_at, target.updated_at diff --git a/internal/storage/queries/kernel_event_authority.sql b/internal/storage/queries/kernel_event_authority.sql index 0ef2a988d..62867aa8d 100644 --- a/internal/storage/queries/kernel_event_authority.sql +++ b/internal/storage/queries/kernel_event_authority.sql @@ -1,53 +1,6 @@ --- name: InsertModelOutputAuthority :one -WITH inserted AS ( - INSERT INTO model_outputs( - agent_id, model_call_context_id, - served_provider_model_slug, stop_reason, provider_replay, created_at - ) - SELECT agent.id, context.id, - sqlc.arg(served_provider_model_slug), sqlc.arg(stop_reason), - sqlc.narg(provider_replay)::jsonb, statement_timestamp() - FROM agents agent - JOIN model_call_contexts context ON context.project_id = agent.project_id - AND context.agent_id = agent.id - AND context.id = sqlc.arg(model_call_context_id)::uuid - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - ON CONFLICT (agent_id, model_call_context_id) DO NOTHING - RETURNING id, agent_id, model_call_context_id, - served_provider_model_slug, stop_reason, provider_replay, created_at -) -SELECT inserted.id, agent.project_id, inserted.agent_id, - ( - SELECT context_turn.turn_id - FROM model_call_context_turns context_turn - WHERE context_turn.project_id = agent.project_id - AND context_turn.agent_id = inserted.agent_id - AND context_turn.model_call_context_id = inserted.model_call_context_id - ) AS turn_id, - inserted.model_call_context_id, - inserted.served_provider_model_slug, inserted.stop_reason, - context.provider_response_id, - inserted.provider_replay, - context.input_tokens_total, context.uncached_input_tokens, - context.cache_read_input_tokens, context.cache_write_input_tokens, - context.output_tokens_total, context.reasoning_output_tokens, - inserted.created_at -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id -JOIN model_call_contexts context ON context.project_id = agent.project_id - AND context.agent_id = inserted.agent_id - AND context.id = inserted.model_call_context_id; - -- name: GetModelOutputByModelContext :one SELECT output.id, agent.project_id, output.agent_id, - ( - SELECT context_turn.turn_id - FROM model_call_context_turns context_turn - WHERE context_turn.project_id = agent.project_id - AND context_turn.agent_id = output.agent_id - AND context_turn.model_call_context_id = output.model_call_context_id - ) AS turn_id, + context.turn_id, output.model_call_context_id, output.served_provider_model_slug, output.stop_reason, context.provider_response_id, @@ -65,162 +18,6 @@ WHERE agent.project_id = sqlc.arg(project_id) AND output.agent_id = sqlc.arg(agent_id) AND output.model_call_context_id = sqlc.arg(model_call_context_id); --- name: InsertControlAgentInput :one -WITH generated AS ( - SELECT uuidv7() AS id -) -INSERT INTO agent_inputs( - id, project_id, agent_id, state, - actor_id, input_kind, delivery_mode, control_type, - idempotency_scope, input_idempotency_key, queued_at, metadata -) -SELECT generated.id, agent.project_id, agent.id, 'received', - sqlc.arg(actor_id), 'control', 'immediate', - sqlc.arg(control_type), sqlc.arg(idempotency_scope), sqlc.arg(input_idempotency_key), - statement_timestamp(), sqlc.arg(metadata) -FROM agents agent -JOIN generated ON true -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) -ON CONFLICT DO NOTHING -RETURNING id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - coalesce(idempotency_scope, '') AS idempotency_scope, - coalesce(input_idempotency_key, '') AS input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, - coalesce(rejected_reason, '') AS rejected_reason, metadata; - --- name: ResolveControlAgentInput :exec -UPDATE agent_inputs -SET state = 'resolved', - admitted_event_id = sqlc.arg(event_id), - admitted_at = statement_timestamp(), - resolved_at = statement_timestamp() -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND id = sqlc.arg(id) - AND input_kind = 'control' - AND control_type = sqlc.arg(control_type) - AND state = 'received'; - --- name: InsertConfigChangeAgentInput :one -WITH generated AS ( - SELECT uuidv7() AS id -), -inserted AS ( -INSERT INTO agent_inputs( - id, project_id, agent_id, state, - actor_id, input_kind, delivery_mode, - agent_config_id, idempotency_scope, input_idempotency_key, queued_at, metadata -) -SELECT generated.id, agent.project_id, agent.id, 'received', - sqlc.narg(actor_id)::uuid, 'config_change', 'immediate', - sqlc.arg(agent_config_id), - sqlc.arg(idempotency_scope), sqlc.arg(input_idempotency_key), - statement_timestamp(), sqlc.arg(metadata) -FROM agents agent -JOIN generated ON true -WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) -ON CONFLICT DO NOTHING -RETURNING id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - coalesce(idempotency_scope, '') AS idempotency_scope, - coalesce(input_idempotency_key, '') AS input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, - coalesce(rejected_reason, '') AS rejected_reason, metadata -) -SELECT id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - idempotency_scope, - input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - control_type, target_interaction_id, agent_config_id, resolved_at, - rejected_reason, metadata -FROM inserted -UNION ALL -SELECT id, project_id, agent_id, - state, input_rank, actor_id, input_kind, - coalesce(idempotency_scope, '') AS idempotency_scope, - coalesce(input_idempotency_key, '') AS input_idempotency_key, - queued_at, admitted_event_id, admitted_at, canceled_at, delivery_mode, - coalesce(control_type, '') AS control_type, target_interaction_id, agent_config_id, resolved_at, - coalesce(rejected_reason, '') AS rejected_reason, metadata -FROM agent_inputs -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND idempotency_scope = sqlc.arg(idempotency_scope) - AND input_idempotency_key = sqlc.arg(input_idempotency_key) -LIMIT 1; - --- name: ResolveConfigChangeAgentInput :one -WITH causal_event AS MATERIALIZED ( - SELECT input.agent_config_id, event.created_at - FROM agent_inputs input - JOIN agent_events event ON event.agent_id = input.agent_id - AND event.agent_input_id = input.id - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.input_kind = 'config_change' - AND input.state = 'received' - AND event.id = sqlc.arg(event_id) - AND event.event_kind = 'agent_input' -), -resolved AS ( - UPDATE agent_inputs input - SET state = 'resolved', - admitted_event_id = sqlc.arg(event_id), - admitted_at = causal_event.created_at, - resolved_at = causal_event.created_at - FROM causal_event - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.arg(id) - AND input.state = 'received' - RETURNING input.admitted_at, input.resolved_at, input.agent_config_id -), -updated_agent AS ( - UPDATE agents agent - SET current_config_id = resolved.agent_config_id, - updated_at = resolved.resolved_at - FROM resolved - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - RETURNING agent.id -) -SELECT resolved.admitted_at, resolved.resolved_at -FROM resolved -JOIN updated_agent ON true; - --- name: InsertToolCallResultAuthority :one -WITH inserted AS ( - INSERT INTO tool_call_results( - agent_id, tool_call_id, outcome, completed_at - ) - SELECT tool_call.agent_id, tool_call.id, - sqlc.arg(outcome), statement_timestamp() - FROM tool_call_read_projection tool_call - WHERE tool_call.project_id = sqlc.arg(project_id) - AND tool_call.agent_id = sqlc.arg(agent_id) - AND tool_call.id = sqlc.arg(tool_call_id) - AND tool_call.state = 'completed' - ON CONFLICT (agent_id, tool_call_id) DO NOTHING - RETURNING id, agent_id, tool_call_id, outcome, completed_at -) -SELECT inserted.id, agent.project_id, inserted.agent_id, - source_event.turn_id, inserted.tool_call_id, inserted.outcome, - inserted.completed_at -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id -JOIN tool_calls tool_call ON tool_call.agent_id = inserted.agent_id - AND tool_call.id = inserted.tool_call_id -JOIN agent_events source_event ON source_event.agent_id = tool_call.agent_id - AND source_event.model_output_id = tool_call.model_output_id - AND source_event.event_kind = 'model_output'; - -- name: GetToolCallResultByToolCall :one SELECT result.id, agent.project_id, result.agent_id, source_event.turn_id AS turn_id, @@ -247,86 +44,6 @@ SELECT EXISTS ( AND event.event_kind = 'tool_result' )::boolean; --- name: ListToolCallResultContentBlocks :many -SELECT block.id, agent.project_id, block.agent_id, block.owner_kind, - block.owner_agent_input_id, block.owner_model_output_id, - block.owner_tool_call_result_id, block.ordinal, block.block_kind, - coalesce(block.text_content, '') AS text_content, block.structured_data, - block.artifact_id, block.tool_call_id, block.exclude_from_model_context, - block.metadata, block.created_at -FROM content_blocks block -JOIN agents agent ON agent.id = block.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND block.agent_id = sqlc.arg(agent_id) - AND block.owner_tool_call_result_id = sqlc.arg(tool_call_result_id) -ORDER BY block.ordinal; - --- name: InsertContentBlock :one -WITH content_owner AS ( - SELECT input.queued_at AS created_at - FROM agent_inputs input - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.narg(owner_agent_input_id)::uuid - - UNION ALL - - SELECT output.created_at - FROM model_outputs output - JOIN agents agent ON agent.id = output.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND output.agent_id = sqlc.arg(agent_id) - AND output.id = sqlc.narg(owner_model_output_id)::uuid - - UNION ALL - - SELECT result.completed_at - FROM tool_call_results result - JOIN agents agent ON agent.id = result.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND result.agent_id = sqlc.arg(agent_id) - AND result.id = sqlc.narg(owner_tool_call_result_id)::uuid -), -inserted AS ( - INSERT INTO content_blocks( - agent_id, owner_kind, owner_agent_input_id, - owner_model_output_id, owner_tool_call_result_id, ordinal, block_kind, - text_content, structured_data, artifact_id, tool_call_id, - exclude_from_model_context, metadata, created_at - ) - SELECT - sqlc.arg(agent_id), sqlc.arg(owner_kind), - sqlc.narg(owner_agent_input_id)::uuid, sqlc.narg(owner_model_output_id)::uuid, - sqlc.narg(owner_tool_call_result_id)::uuid, sqlc.arg(ordinal), sqlc.arg(block_kind), - sqlc.narg(text_content), sqlc.narg(structured_data)::jsonb, sqlc.narg(artifact_id)::uuid, sqlc.narg(tool_call_id)::uuid, - sqlc.arg(exclude_from_model_context), sqlc.arg(metadata)::jsonb, content_owner.created_at - FROM content_owner - RETURNING id, agent_id, owner_kind, owner_agent_input_id, - owner_model_output_id, owner_tool_call_result_id, ordinal, block_kind, - coalesce(text_content, '') AS text_content, structured_data, - artifact_id, tool_call_id, created_at -) -SELECT inserted.id, agent.project_id, inserted.agent_id, inserted.owner_kind, - inserted.owner_agent_input_id, inserted.owner_model_output_id, - inserted.owner_tool_call_result_id, inserted.ordinal, inserted.block_kind, - inserted.text_content, inserted.structured_data, inserted.artifact_id, - inserted.tool_call_id, inserted.created_at -FROM inserted -JOIN agents agent ON agent.id = inserted.agent_id; - --- name: ListContentBlocksForModelOutput :many -SELECT block.id, agent.project_id, block.agent_id, block.owner_kind, - block.owner_agent_input_id, block.owner_model_output_id, - block.owner_tool_call_result_id, block.ordinal, block.block_kind, - coalesce(block.text_content, '') AS text_content, block.structured_data, - block.artifact_id, block.tool_call_id, block.created_at -FROM content_blocks block -JOIN agents agent ON agent.id = block.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND block.agent_id = sqlc.arg(agent_id) - AND block.owner_model_output_id = sqlc.arg(model_output_id) -ORDER BY block.ordinal ASC, block.id ASC; - -- name: ListContentBlocksForAgentInputs :many SELECT block.id, agent.project_id, block.agent_id, block.owner_kind, block.owner_agent_input_id, block.owner_model_output_id, @@ -340,77 +57,6 @@ WHERE agent.project_id = sqlc.arg(project_id) AND block.owner_agent_input_id = ANY(sqlc.arg(agent_input_ids)::uuid[]) ORDER BY block.owner_agent_input_id ASC, block.ordinal ASC, block.id ASC; --- name: GetToolCallProviderCallID :one -SELECT provider_call_id -FROM tool_call_read_projection -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND id = sqlc.arg(id); - --- name: InsertTypedAgentEvent :one -WITH event_time AS ( - SELECT statement_timestamp() AS created_at - FROM agent_inputs input - WHERE input.project_id = sqlc.arg(project_id) - AND input.agent_id = sqlc.arg(agent_id) - AND input.id = sqlc.narg(agent_input_id)::uuid - - UNION ALL - - SELECT output.created_at - FROM model_outputs output - JOIN agents agent ON agent.id = output.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND output.agent_id = sqlc.arg(agent_id) - AND output.id = sqlc.narg(model_output_id)::uuid - - UNION ALL - - SELECT result.completed_at - FROM tool_call_results result - JOIN agents agent ON agent.id = result.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND result.agent_id = sqlc.arg(agent_id) - AND result.id = sqlc.narg(tool_call_result_id)::uuid - - UNION ALL - - SELECT checkpoint.created_at - FROM context_checkpoints checkpoint - JOIN agents agent ON agent.id = checkpoint.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND checkpoint.agent_id = sqlc.arg(agent_id) - AND checkpoint.id = sqlc.narg(context_checkpoint_id)::uuid -) -INSERT INTO agent_events( - id, agent_id, turn_id, sequence, event_kind, - idempotency_key, agent_input_id, model_output_id, - tool_call_result_id, context_checkpoint_id, is_opening_event, created_at -) -SELECT - coalesce(sqlc.narg(id)::uuid, uuidv7()), - sqlc.arg(agent_id), sqlc.arg(turn_id), sqlc.arg(sequence), - sqlc.arg(event_kind), - sqlc.narg(idempotency_key), - sqlc.narg(agent_input_id)::uuid, sqlc.narg(model_output_id)::uuid, - sqlc.narg(tool_call_result_id)::uuid, sqlc.narg(context_checkpoint_id)::uuid, - sqlc.arg(is_opening_event), event_time.created_at -FROM event_time -ON CONFLICT (agent_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING -RETURNING id, agent_id, turn_id, is_opening_event, sequence, event_kind, created_at, coalesce(idempotency_key, '') AS idempotency_key, agent_input_id, model_output_id, tool_call_result_id, context_checkpoint_id; - --- name: GetTypedAgentEventByIdempotency :one -SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, - event.sequence, event.event_kind, event.created_at, - coalesce(event.idempotency_key, '') AS idempotency_key, - event.agent_input_id, event.model_output_id, event.tool_call_result_id, - event.context_checkpoint_id -FROM agent_events event -JOIN agents agent ON agent.id = event.agent_id -WHERE agent.project_id = sqlc.arg(project_id) - AND event.agent_id = sqlc.arg(agent_id) - AND event.idempotency_key = sqlc.arg(idempotency_key)::text; - -- name: GetTypedAgentEventByModelOutput :one SELECT event.id, event.agent_id, event.turn_id, event.is_opening_event, event.sequence, event.event_kind, event.created_at, diff --git a/internal/storage/queries/model_contexts.sql b/internal/storage/queries/model_contexts.sql index 12c02e9bb..3e3f36793 100644 --- a/internal/storage/queries/model_contexts.sql +++ b/internal/storage/queries/model_contexts.sql @@ -1,169 +1,3 @@ --- name: GetModelCallRevisionForClaim :one --- @sqlc-vet-disable configured-models-deleted-at --- @sqlc-vet-disable model-provider-configs-deleted-at --- Model-call lineage must survive soft deletion so the unavailable model can fail durably. -SELECT configured_model.current_revision_id, - CASE - WHEN provider.management_kind = 'cluster' - THEN COALESCE(admission.new_managed_work_allowed, true) - ELSE true - END AS new_managed_work_allowed -FROM agent_configs config -JOIN configured_models configured_model ON configured_model.org_id = config.org_id - AND configured_model.id = config.configured_model_id -JOIN model_provider_configs provider ON provider.org_id = configured_model.org_id - AND provider.id = configured_model.model_provider_config_id -LEFT JOIN org_managed_work_admission admission ON admission.org_id = config.org_id -WHERE config.project_id = sqlc.arg(project_id) - AND config.id = sqlc.arg(agent_config_id); - --- name: InsertNormalModelCallContext :one -WITH active_agent AS MATERIALIZED ( - SELECT agent.org_id, agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - AND agent.state = 'active' -), -live_runtime AS MATERIALIZED ( - SELECT runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN active_agent agent ON agent.id = runtime_lock.agent_id - WHERE runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -), -active_config AS MATERIALIZED ( - SELECT input.agent_config_id - FROM agent_events event - JOIN agent_inputs input ON input.agent_id = event.agent_id - AND input.id = event.agent_input_id - JOIN active_agent agent ON agent.id = event.agent_id - WHERE event.sequence <= sqlc.arg(input_event_sequence) - AND event.event_kind = 'agent_input' - AND input.input_kind = 'config_change' - AND input.state = 'resolved' - AND input.admitted_event_id = event.id - AND input.agent_config_id IS NOT NULL - ORDER BY event.sequence DESC - LIMIT 1 -), -selected_model AS MATERIALIZED ( - SELECT config.id AS agent_config_id, - revision.id AS configured_model_revision_id - FROM agent_configs config - JOIN active_config ON active_config.agent_config_id = config.id - JOIN configured_model_revisions revision ON revision.org_id = config.org_id - AND revision.configured_model_id = config.configured_model_id - AND revision.id = sqlc.arg(configured_model_revision_id) - WHERE config.project_id = sqlc.arg(project_id) -) -INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, attempt_number, - agent_config_id, configured_model_revision_id, input_event_sequence, - runtime_lock_id, state, created_at -) -SELECT active_agent.org_id, - active_agent.project_id, - active_agent.id, - 'normal', - 1, - selected_model.agent_config_id, - selected_model.configured_model_revision_id, - sqlc.arg(input_event_sequence), - live_runtime.id, - 'started', - statement_timestamp() -FROM active_agent -JOIN live_runtime ON true -JOIN selected_model ON true -RETURNING id; - --- name: InsertTriggeredCompactionModelCallContext :one -WITH active_agent AS MATERIALIZED ( - SELECT agent.org_id, agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - AND agent.state = 'active' -), -live_runtime AS MATERIALIZED ( - SELECT runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN active_agent agent ON agent.id = runtime_lock.agent_id - WHERE runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -), -latest_compaction AS MATERIALIZED ( - SELECT dependency.id, dependency.state, dependency.recovery_kind, - dependency.input_event_sequence, dependency.source_event_sequence_end - FROM model_call_contexts dependency - JOIN model_call_contexts parent ON parent.project_id = dependency.project_id - AND parent.agent_id = dependency.agent_id - AND parent.id = sqlc.arg(parent_model_call_context_id) - JOIN active_agent agent ON agent.project_id = parent.project_id - AND agent.id = parent.agent_id - WHERE dependency.operation_kind = 'compaction' - AND dependency.input_event_sequence = parent.input_event_sequence - ORDER BY dependency.source_event_sequence_end ASC, - dependency.attempt_number DESC - LIMIT 1 -), -blocked AS MATERIALIZED ( - SELECT context.id, context.agent_config_id, - context.input_event_sequence - FROM model_call_contexts context - JOIN active_agent agent ON agent.project_id = context.project_id - AND agent.id = context.agent_id - LEFT JOIN latest_compaction latest ON true - WHERE context.id = sqlc.arg(parent_model_call_context_id) - AND context.operation_kind = 'normal' - AND context.state = 'failed' - AND context.recovery_kind = 'compact' - AND ( - latest.id IS NULL - OR ( - latest.state = 'failed' - AND latest.recovery_kind = 'reduce_compaction_source' - AND latest.input_event_sequence = context.input_event_sequence - AND sqlc.arg(source_event_sequence_end)::bigint < latest.source_event_sequence_end - ) - ) -), -selected_model AS MATERIALIZED ( - SELECT blocked.id AS blocked_context_id, - revision.id AS configured_model_revision_id - FROM blocked - JOIN agent_configs config ON config.project_id = sqlc.arg(project_id) - AND config.id = blocked.agent_config_id - JOIN configured_model_revisions revision ON revision.org_id = config.org_id - AND revision.configured_model_id = config.configured_model_id - AND revision.id = sqlc.arg(configured_model_revision_id) -) -INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, attempt_number, - agent_config_id, configured_model_revision_id, input_event_sequence, - source_event_sequence_end, runtime_lock_id, state, created_at -) -SELECT active_agent.org_id, - active_agent.project_id, - active_agent.id, - 'compaction', - 1, - blocked.agent_config_id, - selected_model.configured_model_revision_id, - blocked.input_event_sequence, - sqlc.arg(source_event_sequence_end), - live_runtime.id, - 'started', - statement_timestamp() -FROM active_agent -JOIN live_runtime ON true -JOIN blocked ON true -JOIN selected_model ON selected_model.blocked_context_id = blocked.id -RETURNING id; - -- name: GetNormalModelCallContextByIdentity :one SELECT context.id FROM model_call_contexts context @@ -174,17 +8,6 @@ WHERE context.project_id = sqlc.arg(project_id) ORDER BY context.attempt_number DESC LIMIT 1; --- name: GetCompactionModelCallContextByIdentity :one -SELECT context.id -FROM model_call_contexts context -WHERE context.project_id = sqlc.arg(project_id) - AND context.agent_id = sqlc.arg(agent_id) - AND context.operation_kind = 'compaction' - AND context.input_event_sequence = sqlc.arg(input_event_sequence) - AND context.source_event_sequence_end = sqlc.arg(source_event_sequence_end) -ORDER BY context.attempt_number DESC -LIMIT 1; - -- name: GetModelCallContext :one SELECT context.id, context.org_id, context.project_id, context.agent_id, context.operation_kind, context.attempt_number, context.agent_config_id, @@ -205,13 +28,6 @@ WHERE context.project_id = sqlc.arg(project_id) AND context.agent_id = sqlc.arg(agent_id) AND context.id = sqlc.arg(id); --- name: ModelCallContextHasLaterSemanticEvent :one -SELECT model_call_context_has_later_semantic_event( - sqlc.arg(project_id), - sqlc.arg(agent_id), - sqlc.arg(model_call_context_id) -) AS has_later_semantic_event; - -- name: GetProviderReplaySuppressionCutoff :one -- @sqlc-vet-disable model-provider-configs-deleted-at -- Replay safety remains part of immutable call lineage after provider soft deletion. @@ -246,248 +62,3 @@ SELECT COALESCE(( AND failure.api_variant = current.api_variant ), 0)::bigint AS cutoff_event_sequence FROM current_replay_identity current; - --- name: GetModelCallContextTurnID :one -SELECT turn_id -FROM model_call_context_turns -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND model_call_context_id = sqlc.arg(model_call_context_id); - --- name: GetModelCallContextProviderModelSlug :one -SELECT revision.provider_model_slug -FROM model_call_contexts context -JOIN configured_model_revisions revision ON revision.org_id = context.org_id - AND revision.id = context.configured_model_revision_id -WHERE context.project_id = sqlc.arg(project_id) - AND context.agent_id = sqlc.arg(agent_id) - AND context.id = sqlc.arg(model_call_context_id); - --- name: InsertNextModelCallContext :one -WITH agent_scope AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) -), -predecessor AS MATERIALIZED ( - SELECT context.id, context.org_id, context.project_id, context.agent_id, - context.operation_kind, context.attempt_number, - context.agent_config_id, context.input_event_sequence, - context.source_event_sequence_end - FROM model_call_contexts context - JOIN agent_scope agent ON agent.project_id = context.project_id - AND agent.id = context.agent_id - WHERE context.id = sqlc.arg(predecessor_model_call_context_id) - AND context.state = 'failed' - AND context.recovery_kind = 'retry' - AND context.retry_at <= statement_timestamp() - AND context.attempt_number <= sqlc.arg(max_retries)::integer - AND EXISTS ( - SELECT 1 - FROM agent_continuable_model_contexts(context.project_id, context.agent_id) continuable - WHERE continuable.model_call_context_id = context.id - AND NOT continuable.has_later_semantic_event - ) -), -selected_model AS MATERIALIZED ( - SELECT predecessor.id AS predecessor_context_id, - revision.id AS configured_model_revision_id - FROM predecessor - JOIN agent_configs config ON config.project_id = predecessor.project_id - AND config.id = predecessor.agent_config_id - JOIN configured_model_revisions revision ON revision.org_id = config.org_id - AND revision.configured_model_id = config.configured_model_id - AND revision.id = sqlc.arg(configured_model_revision_id) -), -live_runtime AS MATERIALIZED ( - SELECT runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN predecessor context ON context.agent_id = runtime_lock.agent_id - WHERE runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -INSERT INTO model_call_contexts( - org_id, project_id, agent_id, operation_kind, attempt_number, - agent_config_id, configured_model_revision_id, input_event_sequence, - source_event_sequence_end, runtime_lock_id, state, created_at -) -SELECT predecessor.org_id, - predecessor.project_id, - predecessor.agent_id, - predecessor.operation_kind, - predecessor.attempt_number + 1, - predecessor.agent_config_id, - selected_model.configured_model_revision_id, - predecessor.input_event_sequence, - predecessor.source_event_sequence_end, - live_runtime.id, - 'started', - statement_timestamp() -FROM predecessor -JOIN selected_model ON selected_model.predecessor_context_id = predecessor.id -JOIN live_runtime ON true -RETURNING id; - --- name: GetLiveModelCallContextForRuntime :one -SELECT id -FROM model_call_contexts -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND runtime_lock_id = sqlc.arg(runtime_lock_id) - AND state = 'started'; - --- name: FinishModelCallContext :one -WITH agent_scope AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) -), -runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agent_scope agent ON agent.id = runtime_lock.agent_id - WHERE runtime_lock.id = sqlc.arg(runtime_lock_id) - AND ( - sqlc.arg(allow_inactive_runtime_lock_for_teardown)::boolean - OR ( - runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() - ) - ) -) -UPDATE model_call_contexts context -SET state = sqlc.arg(to_state), - recovery_kind = sqlc.narg(recovery_kind), - api_format = sqlc.arg(api_format), - api_variant = sqlc.arg(api_variant), - provider_request_id = sqlc.arg(provider_request_id), - provider_response_id = sqlc.arg(provider_response_id), - error_kind = sqlc.arg(error_kind), - error_code = sqlc.arg(error_code), - error_message = sqlc.arg(error_message), - error_details = sqlc.arg(error_details), - retry_at = CASE - WHEN sqlc.narg(retry_delay_microseconds)::bigint IS NULL THEN NULL - ELSE statement_timestamp() + - (sqlc.narg(retry_delay_microseconds)::bigint * interval '1 microsecond') - END, - input_tokens_total = sqlc.narg(input_tokens_total)::integer, - uncached_input_tokens = sqlc.narg(uncached_input_tokens)::integer, - cache_read_input_tokens = sqlc.narg(cache_read_input_tokens)::integer, - cache_write_input_tokens = sqlc.narg(cache_write_input_tokens)::integer, - output_tokens_total = sqlc.narg(output_tokens_total)::integer, - reasoning_output_tokens = sqlc.narg(reasoning_output_tokens)::integer, - provider_reported_cost_usd = sqlc.narg(provider_reported_cost_usd)::text::numeric, - provider_metadata = sqlc.arg(provider_metadata), - completed_at = statement_timestamp() -FROM runtime -WHERE context.project_id = runtime.project_id - AND context.agent_id = runtime.agent_id - AND context.id = sqlc.arg(id) - AND context.runtime_lock_id = runtime.id - AND context.state = 'started' -RETURNING context.id; - --- name: AgentHasLiveModelCallContextBeforeFrontier :one -SELECT EXISTS ( - SELECT 1 - FROM model_call_contexts context - WHERE context.project_id = sqlc.arg(project_id) - AND context.agent_id = sqlc.arg(agent_id) - AND context.state = 'started' - AND context.input_event_sequence < sqlc.arg(input_event_sequence) -)::boolean; - --- name: AgentHasLiveModelCallContexts :one -SELECT EXISTS ( - SELECT 1 - FROM model_call_contexts context - WHERE context.project_id = sqlc.arg(project_id) - AND context.agent_id = sqlc.arg(agent_id) - AND context.state = 'started' -)::boolean; - --- name: CancelRuntimeModelCallContextsForLifecycle :execrows -UPDATE model_call_contexts context -SET state = 'canceled', - recovery_kind = NULL, - error_kind = 'canceled', - error_code = sqlc.arg(error_code), - error_message = sqlc.arg(error_message), - error_details = sqlc.arg(error_details), - retry_at = NULL, - completed_at = statement_timestamp() -WHERE context.project_id = sqlc.arg(project_id) - AND context.agent_id = sqlc.arg(agent_id) - AND context.runtime_lock_id = sqlc.arg(runtime_lock_id) - AND context.state = 'started'; - --- name: InterruptRuntimeModelCallContextForRetry :execrows -UPDATE model_call_contexts context -SET state = 'failed', - recovery_kind = 'retry', - error_kind = sqlc.arg(error_kind), - error_code = sqlc.arg(error_code), - error_message = sqlc.arg(error_message), - error_details = sqlc.arg(error_details)::jsonb, - retry_at = statement_timestamp() + - (sqlc.arg(retry_delay_microseconds)::bigint * interval '1 microsecond'), - completed_at = statement_timestamp() -WHERE context.project_id = sqlc.arg(project_id) - AND context.agent_id = sqlc.arg(agent_id) - AND context.id = sqlc.arg(id) - AND context.runtime_lock_id = sqlc.arg(runtime_lock_id) - AND context.state = 'started'; - --- name: OpeningContentInputSetMatchesInputSequence :one -WITH requested AS ( - SELECT input_id - FROM unnest(sqlc.arg(input_ids)::uuid[]) AS input_ids(input_id) -), -requested_distinct AS ( - SELECT DISTINCT input_id - FROM requested -), -target_turn AS MATERIALIZED ( - SELECT agent_turn_id_at_event_sequence( - sqlc.arg(agent_id), - sqlc.arg(input_event_sequence) - ) AS turn_id -), -opening AS ( - SELECT opening.input_id - FROM target_turn turn - CROSS JOIN LATERAL agent_model_call_opening_content_inputs( - sqlc.arg(project_id), - sqlc.arg(agent_id), - turn.turn_id, - sqlc.arg(input_event_sequence) - ) AS opening(input_id, event_sequence) -) -SELECT ( - (SELECT count(*) FROM requested) = (SELECT count(*) FROM requested_distinct) - AND NOT EXISTS ( - SELECT input_id FROM requested_distinct - EXCEPT - SELECT input_id FROM opening - ) - AND NOT EXISTS ( - SELECT input_id FROM opening - EXCEPT - SELECT input_id FROM requested_distinct - ) -)::boolean; - --- name: ListModelCallOpeningContentInputs :many -SELECT opening.input_id::uuid AS input_id, - opening.event_sequence::bigint AS event_sequence -FROM agent_model_call_opening_content_inputs( - sqlc.arg(project_id), - sqlc.arg(agent_id), - sqlc.arg(turn_id), - sqlc.arg(input_event_sequence) -) AS opening(input_id, event_sequence) -ORDER BY opening.event_sequence; diff --git a/internal/storage/queries/processes.sql b/internal/storage/queries/processes.sql index 5bbfc5360..d04042bfb 100644 --- a/internal/storage/queries/processes.sql +++ b/internal/storage/queries/processes.sql @@ -76,19 +76,6 @@ JOIN reachable_machine ON true ON CONFLICT (agent_id, tool_call_id) DO NOTHING RETURNING id, org_id, project_id, agent_id, tool_call_id, runtime_lock_id, agent_machine_binding_id, machine_id, execution_granted_at, cwd, env, secret_env, timeout_seconds, initial_wait_ms, default_output_cursor, state, state_reason_code, state_reason_message, source_started_at, source_ended_at, state_changed_at, exit_code, exit_signal, created_at, updated_at, last_activity_at, execution_spec; --- name: MachineReachableForProjectMachine :one -SELECT machine.id -FROM machines machine -JOIN project_machine_grants pmgrant ON pmgrant.org_id = machine.org_id - AND pmgrant.machine_id = machine.id - AND pmgrant.project_id = sqlc.arg(project_id) -JOIN machine_connection_states connection ON connection.org_id = machine.org_id - AND connection.machine_id = machine.id -WHERE machine.id = sqlc.arg(machine_id) - AND machine.deleted_at IS NULL - AND connection.connection_state IN ('online', 'asleep') -LIMIT 1; - -- name: CountNonTerminalProcessesForAgent :one SELECT count(*)::bigint FROM processes @@ -459,7 +446,7 @@ WHERE org_id = sqlc.arg(org_id) AND state IN ('starting', 'running') RETURNING id, org_id, project_id, agent_id, tool_call_id, runtime_lock_id, agent_machine_binding_id, machine_id, execution_granted_at, cwd, env, secret_env, timeout_seconds, initial_wait_ms, default_output_cursor, state, state_reason_code, state_reason_message, source_started_at, source_ended_at, state_changed_at, exit_code, exit_signal, created_at, updated_at, last_activity_at, execution_spec; --- name: LockAgentsForExecutionRevoked :many +-- name: ListAgentsForExecutionRevoked :many SELECT agent.id FROM agents agent WHERE agent.project_id = sqlc.arg(project_id) @@ -490,8 +477,7 @@ WHERE agent.project_id = sqlc.arg(project_id) ) ) ) -ORDER BY agent.id -FOR UPDATE; +ORDER BY agent.id; -- name: ListProcessesForExecutionRevoked :many SELECT process.id, process.org_id, process.project_id, process.agent_id, process.tool_call_id, process.runtime_lock_id, process.agent_machine_binding_id, process.machine_id, process.execution_granted_at, process.cwd, process.env, process.secret_env, process.timeout_seconds, process.initial_wait_ms, process.default_output_cursor, process.state, process.state_reason_code, process.state_reason_message, process.source_started_at, process.source_ended_at, process.state_changed_at, process.exit_code, process.exit_signal, process.created_at, process.updated_at, process.last_activity_at, process.execution_spec @@ -598,3 +584,16 @@ WHERE project_id = sqlc.arg(project_id) AND state IN ('starting', 'running') ORDER BY updated_at DESC, id LIMIT 20; + +-- name: MachineReachableForProjectMachine :one +SELECT machine.id +FROM machines machine +JOIN project_machine_grants pmgrant ON pmgrant.org_id = machine.org_id + AND pmgrant.machine_id = machine.id + AND pmgrant.project_id = sqlc.arg(project_id) +JOIN machine_connection_states connection ON connection.org_id = machine.org_id + AND connection.machine_id = machine.id +WHERE machine.id = sqlc.arg(machine_id) + AND machine.deleted_at IS NULL + AND connection.connection_state IN ('online', 'asleep') +LIMIT 1; diff --git a/internal/storage/queries/secrets.sql b/internal/storage/queries/secrets.sql index 9242f757a..90a6d0e91 100644 --- a/internal/storage/queries/secrets.sql +++ b/internal/storage/queries/secrets.sql @@ -65,6 +65,16 @@ JOIN secret_versions v AND v.id = s.current_version_id WHERE s.org_id = sqlc.arg(org_id) AND s.id = sqlc.arg(id) AND s.deleted_at IS NULL; +-- name: GetCurrentSecretPayload :one +SELECT v.id, v.org_id, v.secret_id, v.version_number, v.payload_keys, v.encryption_scheme, + v.key_id, v.dek_wrapped_by, v.encrypted_dek, v.encrypted_dek_nonce, v.nonce, v.ciphertext, v.created_at, + (v.oauth_access_token_expires_at IS NOT NULL)::boolean AS oauth_access_token_expires, + coalesce(greatest(extract(epoch FROM v.oauth_access_token_expires_at - statement_timestamp()), 0), 0)::double precision AS oauth_access_token_remaining_seconds, + s.owner_kind, s.management_kind, s.kind +FROM secrets s +JOIN secret_versions v ON v.org_id = s.org_id AND v.secret_id = s.id AND v.id = s.current_version_id +WHERE s.org_id = sqlc.arg(org_id) AND s.id = sqlc.arg(id) AND s.deleted_at IS NULL; + -- name: GetSecretByOwnerName :one SELECT s.id, s.org_id, s.management_kind, s.owner_kind, s.owner_project_id, s.owner_user_id, s.name, s.kind, s.metadata, s.current_version_id, v.version_number AS current_version_number, v.payload_keys, s.created_at, s.updated_at FROM secrets s diff --git a/internal/storage/queries/subagents.sql b/internal/storage/queries/subagents.sql index a9676569e..04f6f8769 100644 --- a/internal/storage/queries/subagents.sql +++ b/internal/storage/queries/subagents.sql @@ -35,8 +35,8 @@ SELECT agent.id, FROM agent_wakeups wake WHERE wake.agent_id = agent.id ) - OR agent_next_wakeup_ready_at(agent.project_id, agent.id) IS NOT NULL - OR agent_has_incomplete_tool_batch(agent.project_id, agent.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.logical_ready_at IS NOT NULL) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.incomplete_tools) )::boolean AS is_running FROM agents agent WHERE agent.project_id = sqlc.arg(project_id) @@ -79,8 +79,8 @@ SELECT agent.id, FROM agent_wakeups wake WHERE wake.agent_id = agent.id ) - OR agent_next_wakeup_ready_at(agent.project_id, agent.id) IS NOT NULL - OR agent_has_incomplete_tool_batch(agent.project_id, agent.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.logical_ready_at IS NOT NULL) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=agent.id AND h.incomplete_tools) )::boolean AS is_running FROM agents agent WHERE agent.project_id = ANY(sqlc.arg(project_ids)::uuid[]) @@ -223,7 +223,7 @@ WHERE NOT EXISTS ( FROM agent_wakeups wake WHERE wake.agent_id = subtree.id ) - OR agent_has_incomplete_tool_batch(subtree.project_id, subtree.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=subtree.id AND h.incomplete_tools) ) ) ORDER BY candidate.created_at, candidate.id @@ -271,6 +271,6 @@ SELECT EXISTS ( FROM agent_wakeups wake WHERE wake.agent_id = subtree.id ) - OR agent_has_incomplete_tool_batch(subtree.project_id, subtree.id) + OR EXISTS (SELECT 1 FROM agent_execution_state h WHERE h.agent_id=subtree.id AND h.incomplete_tools) ) )::boolean AS idle; diff --git a/internal/storage/queries/tool_calls.sql b/internal/storage/queries/tool_calls.sql index 760ea6c0a..684819b54 100644 --- a/internal/storage/queries/tool_calls.sql +++ b/internal/storage/queries/tool_calls.sql @@ -1,77 +1,13 @@ --- name: InsertToolCall :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -), -source_event AS MATERIALIZED ( - SELECT event.id, event.turn_id, event.model_output_id, model_output.created_at - FROM agent_events event - JOIN agents agent ON agent.id = event.agent_id - JOIN model_outputs model_output ON model_output.agent_id = event.agent_id - AND model_output.id = event.model_output_id - WHERE agent.project_id = sqlc.arg(project_id) - AND event.agent_id = sqlc.arg(agent_id) - AND event.id = sqlc.arg(source_event_id) - AND event.event_kind = 'model_output' - AND event.model_output_id IS NOT NULL - AND model_output.model_call_context_id = sqlc.arg(model_call_context_id) -), -candidate AS MATERIALIZED ( - SELECT coalesce(sqlc.narg(tool_call_id)::uuid, uuidv7()) AS id, - mcc.project_id, - mcc.agent_id, - mcc.id AS model_call_context_id, - source_event.id AS source_event_id, - source_event.turn_id, - source_event.model_output_id, - sqlc.arg(provider_call_id)::text AS provider_call_id, - sqlc.arg(name)::text AS name, - sqlc.arg(input)::jsonb AS input, - sqlc.arg(type)::text AS type, - source_event.created_at - FROM model_call_contexts mcc - JOIN live_runtime runtime_lock ON runtime_lock.project_id = mcc.project_id - AND runtime_lock.agent_id = mcc.agent_id - AND runtime_lock.id = mcc.runtime_lock_id - JOIN source_event ON true - WHERE mcc.project_id = sqlc.arg(project_id) - AND mcc.agent_id = sqlc.arg(agent_id) - AND mcc.id = sqlc.arg(model_call_context_id) - AND mcc.runtime_lock_id = sqlc.arg(runtime_lock_id) -), -inserted_call AS ( - INSERT INTO tool_calls( - id, agent_id, model_output_id, - provider_call_id, name, input, - type, state, created_at - ) - SELECT id, agent_id, model_output_id, - provider_call_id, name, input, - type, 'awaiting_authorization', created_at - FROM candidate - RETURNING id, agent_id, - model_output_id, provider_call_id, name, - input, type, state, - runtime_lock_id, created_at -) -SELECT call.id, candidate.project_id, call.agent_id, - candidate.turn_id, candidate.source_event_id, - candidate.model_call_context_id, call.provider_call_id, - call.name, call.input, - call.type, call.state, - ''::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at, NULL::timestamptz AS completed_at -FROM inserted_call call -JOIN candidate ON candidate.id = call.id; - -- name: GetToolCall :one +WITH selected_call AS MATERIALIZED ( + SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, + tc.model_call_context_id, tc.model_output_id, tc.provider_call_id, tc.name, tc.input, + tc.type, tc.state, tc.runtime_lock_id, tc.created_at + FROM tool_call_read_projection tc + WHERE tc.project_id = sqlc.arg(project_id) + AND tc.agent_id = sqlc.arg(agent_id) + AND tc.id = sqlc.arg(id) +) SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.model_call_context_id, tc.provider_call_id, tc.name, tc.input, tc.type, @@ -80,7 +16,7 @@ SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, tc.created_at, result.completed_at -FROM tool_call_read_projection tc +FROM selected_call tc LEFT JOIN tool_call_results result ON result.agent_id = tc.agent_id AND result.tool_call_id = tc.id LEFT JOIN LATERAL ( @@ -100,10 +36,19 @@ LEFT JOIN LATERAL ( FROM content_blocks block WHERE block.agent_id = result.agent_id AND block.owner_tool_call_result_id = result.id -) result_blocks ON true -WHERE tc.project_id = sqlc.arg(project_id) AND tc.agent_id = sqlc.arg(agent_id) AND tc.id = sqlc.arg(id); +) result_blocks ON true; -- name: GetToolCallByProviderCall :one +WITH selected_call AS MATERIALIZED ( + SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, + tc.model_call_context_id, tc.model_output_id, tc.provider_call_id, tc.name, tc.input, + tc.type, tc.state, tc.runtime_lock_id, tc.created_at + FROM tool_call_read_projection tc + WHERE tc.project_id = sqlc.arg(project_id) + AND tc.agent_id = sqlc.arg(agent_id) + AND tc.model_call_context_id = sqlc.arg(model_call_context_id) + AND tc.provider_call_id = sqlc.arg(provider_call_id) +) SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.model_call_context_id, tc.provider_call_id, tc.name, tc.input, tc.type, @@ -112,7 +57,7 @@ SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, tc.created_at, result.completed_at -FROM tool_call_read_projection tc +FROM selected_call tc LEFT JOIN tool_call_results result ON result.agent_id = tc.agent_id AND result.tool_call_id = tc.id LEFT JOIN LATERAL ( @@ -132,45 +77,7 @@ LEFT JOIN LATERAL ( FROM content_blocks block WHERE block.agent_id = result.agent_id AND block.owner_tool_call_result_id = result.id -) result_blocks ON true -WHERE tc.project_id = sqlc.arg(project_id) AND tc.agent_id = sqlc.arg(agent_id) AND tc.model_call_context_id = sqlc.arg(model_call_context_id) AND tc.provider_call_id = sqlc.arg(provider_call_id); - --- name: ListToolCallsForModelContext :many -SELECT tc.id, tc.project_id, tc.agent_id, - tc.turn_id, tc.source_event_id, tc.model_call_context_id, - tc.provider_call_id, tc.name, tc.input, tc.type, tc.state, - coalesce(result.outcome, '') AS outcome, tc.runtime_lock_id, - coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, - tc.created_at, result.completed_at -FROM tool_call_read_projection tc -LEFT JOIN tool_call_results result - ON result.agent_id = tc.agent_id - AND result.tool_call_id = tc.id -LEFT JOIN LATERAL ( - SELECT coalesce(jsonb_agg( - CASE - WHEN block.block_kind = 'text' THEN jsonb_build_object('type', 'text', 'text', block.text_content) - WHEN block.block_kind = 'structured_data' THEN jsonb_build_object('type', 'structured_data', 'value', block.structured_data) - WHEN block.block_kind = 'artifact' THEN - jsonb_build_object('type', 'media_ref', 'artifact_id', block.artifact_id::text) - || CASE - WHEN block.exclude_from_model_context THEN jsonb_build_object('exclude_from_model_context', true) - ELSE '{}'::jsonb - END - END - ORDER BY block.ordinal, block.id - ) FILTER ( - WHERE block.id IS NOT NULL - AND block.block_kind IN ('text', 'structured_data', 'artifact') - ), '[]'::jsonb) AS content_parts - FROM content_blocks block - WHERE block.agent_id = result.agent_id - AND block.owner_tool_call_result_id = result.id -) result_blocks ON true -WHERE tc.project_id = sqlc.arg(project_id) - AND tc.agent_id = sqlc.arg(agent_id) - AND tc.model_call_context_id = sqlc.arg(model_call_context_id) -ORDER BY tc.created_at, tc.id; +) result_blocks ON true; -- name: ListToolCallsForAgent :many SELECT call.id, call.project_id, call.agent_id, @@ -225,468 +132,18 @@ WHERE call.project_id = sqlc.arg(project_id) ORDER BY coalesce(call_block.ordinal, 2147483647), call.created_at, call.id LIMIT 1; --- name: StartToolCall :execrows -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = CASE - WHEN sqlc.arg(retain_runtime_ownership)::boolean THEN 'running' - ELSE 'waiting' - END, - runtime_lock_id = CASE - WHEN sqlc.arg(retain_runtime_ownership)::boolean THEN runtime_lock.id - ELSE NULL - END -FROM live_runtime runtime_lock -WHERE call.agent_id = sqlc.arg(agent_id) - AND call.id = sqlc.arg(id) - AND call.state = 'ready' - AND call.runtime_lock_id IS NULL - AND call.type IN ('built_in', 'mcp') - AND ( - sqlc.arg(retain_runtime_ownership)::boolean - OR call.type = 'built_in' - ) - AND runtime_lock.agent_id = call.agent_id - AND runtime_lock.id = sqlc.arg(runtime_lock_id); - --- name: GetToolCallDispatchState :one -SELECT state, runtime_lock_id -FROM tool_call_read_projection -WHERE project_id = sqlc.arg(project_id) - AND agent_id = sqlc.arg(agent_id) - AND id = sqlc.arg(id); - --- name: RequeueRuntimeToolCall :execrows -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'ready', - runtime_lock_id = NULL -FROM live_runtime runtime_lock -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = sqlc.arg(id) - AND call.state = 'running' - AND call.runtime_lock_id = runtime_lock.id; - --- name: MarkToolCallAwaitingPermission :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'awaiting_permission' -FROM live_runtime runtime_lock -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = sqlc.arg(id) - AND call.state = 'awaiting_authorization' - AND runtime_lock.id = sqlc.arg(runtime_lock_id) -RETURNING call.id; - --- name: MarkToolCallReady :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'ready' -FROM live_runtime runtime_lock -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = sqlc.arg(id) - AND call.state = 'awaiting_authorization' - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND projection.project_id = runtime_lock.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - ''::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at, NULL::timestamptz AS completed_at; - --- name: MarkToolCallReadyFromInteraction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -), -permission_interaction AS MATERIALIZED ( - SELECT agent.project_id, interaction.agent_id, interaction.tool_call_id - FROM agent_interactions interaction - JOIN locked_agent agent ON agent.id = interaction.agent_id - WHERE interaction.agent_id = sqlc.arg(agent_id) - AND interaction.id = sqlc.arg(interaction_id) - AND interaction.interaction_kind = 'permission' - AND interaction.state = 'resolved' -) -UPDATE tool_calls call -SET state = 'ready' -FROM permission_interaction interaction -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = interaction.agent_id - AND call.id = interaction.tool_call_id - AND call.state = 'awaiting_permission' - AND projection.project_id = interaction.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - ''::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at, NULL::timestamptz AS completed_at; - --- name: CompleteToolCallFromPermissionInteraction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -), -permission_interaction AS MATERIALIZED ( - SELECT agent.project_id, interaction.agent_id, interaction.tool_call_id - FROM agent_interactions interaction - JOIN locked_agent agent ON agent.id = interaction.agent_id - WHERE interaction.agent_id = sqlc.arg(agent_id) - AND interaction.id = sqlc.arg(interaction_id) - AND interaction.interaction_kind = 'permission' - AND interaction.state IN ('resolved', 'canceled') -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM permission_interaction interaction -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = interaction.agent_id - AND call.id = interaction.tool_call_id - AND call.state = 'awaiting_permission' - AND projection.project_id = interaction.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteToolCall :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM live_runtime runtime_lock -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = sqlc.arg(agent_id) - AND call.id = sqlc.arg(id) - AND call.state <> 'completed' - AND ( - (call.type IN ('built_in', 'mcp') AND call.state = 'ready') - OR ( - call.state IN ('awaiting_authorization', 'awaiting_permission') - AND sqlc.arg(outcome)::text IN ('failed', 'denied', 'canceled') - ) - ) - AND runtime_lock.agent_id = call.agent_id - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND projection.project_id = runtime_lock.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteRuntimeToolCall :one -WITH live_runtime AS MATERIALIZED ( - SELECT agent.project_id, runtime_lock.agent_id, runtime_lock.id - FROM agent_runtime_locks runtime_lock - JOIN agents agent ON agent.id = runtime_lock.agent_id - WHERE agent.project_id = sqlc.arg(project_id) - AND runtime_lock.agent_id = sqlc.arg(agent_id) - AND runtime_lock.id = sqlc.arg(runtime_lock_id) - AND runtime_lock.cancel_requested_at IS NULL - AND runtime_lock.lease_expires_at > statement_timestamp() -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM live_runtime runtime_lock -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = runtime_lock.agent_id - AND call.id = sqlc.arg(id) - AND call.state = 'running' - AND call.runtime_lock_id = runtime_lock.id - AND call.type IN ('built_in', 'mcp') - AND projection.project_id = runtime_lock.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteWaitingBuiltInToolCall :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = sqlc.arg(id) - AND call.state = 'waiting' - AND call.type = 'built_in' - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteCustomToolCall :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - JOIN tool_calls tool_call ON tool_call.agent_id = agent.id - WHERE agent.project_id = sqlc.arg(project_id) - AND tool_call.agent_id = sqlc.arg(agent_id) - AND tool_call.id = sqlc.arg(id) - FOR UPDATE OF agent -) -UPDATE tool_calls call -SET state = 'completed' -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = sqlc.arg(id) - AND call.state = 'ready' - AND call.runtime_lock_id IS NULL - AND call.type = 'custom' - AND sqlc.arg(outcome)::text IN ('succeeded', 'failed') - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteToolCallFromProcess :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM processes tp -CROSS JOIN locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = sqlc.arg(id) - AND call.state = 'waiting' - AND call.type = 'built_in' - AND tp.project_id = agent.project_id - AND tp.agent_id = call.agent_id - AND tp.id = sqlc.arg(process_id) - AND tp.tool_call_id = call.id - AND tp.state IN ('exited', 'failed', 'killed', 'unknown') - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteToolCallFromQuestionInteraction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.state IN ('running', 'waiting') - AND call.type = 'built_in' - AND EXISTS ( - SELECT 1 - FROM agent_interactions question - WHERE question.agent_id = call.agent_id - AND question.tool_call_id = call.id - AND question.id = sqlc.arg(interaction_id) - AND question.interaction_kind = 'question' - AND question.state IN ('resolved', 'canceled') - ) - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteToolCallFromStartedProcess :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM processes process -CROSS JOIN locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = sqlc.arg(id) - AND call.state = 'waiting' - AND call.type = 'built_in' - AND process.project_id = agent.project_id - AND process.agent_id = call.agent_id - AND process.id = sqlc.arg(process_id) - AND process.tool_call_id = call.id - AND process.state = 'running' - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - --- name: CompleteToolCallFromProcessAction :one -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM process_actions action -CROSS JOIN locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.id = action.tool_call_id - AND call.id = sqlc.arg(tool_call_id) - AND call.state = 'waiting' - AND call.type = 'built_in' - AND action.project_id = agent.project_id - AND action.agent_id = call.agent_id - AND action.process_id = sqlc.arg(process_id) - AND action.id = sqlc.arg(process_action_id) - AND action.state IN ('applied', 'failed', 'unknown') - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - sqlc.arg(outcome)::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; - -- name: ListCompletedToolCallsAtWatermark :many +WITH selected_calls AS MATERIALIZED ( + SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, + tc.model_call_context_id, tc.model_output_id, tc.provider_call_id, tc.name, tc.input, + tc.type, tc.state, tc.runtime_lock_id, tc.created_at + FROM tool_call_read_projection tc + WHERE tc.project_id = sqlc.arg(project_id) + AND tc.agent_id = sqlc.arg(agent_id) + AND tc.state = 'completed' + AND tc.source_event_sequence > sqlc.arg(after_event_sequence) + AND tc.source_event_sequence <= sqlc.arg(max_event_sequence) +) SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.source_event_id, tc.source_event_sequence, tc.model_call_context_id, result.id AS tool_call_result_id, @@ -696,7 +153,7 @@ SELECT tc.id, tc.project_id, tc.agent_id, tc.turn_id, tc.runtime_lock_id, coalesce(result_blocks.content_parts, '[]'::jsonb)::jsonb AS result_content_parts, tc.created_at, result.completed_at -FROM tool_call_read_projection tc +FROM selected_calls tc JOIN content_blocks call_block ON call_block.agent_id = tc.agent_id AND call_block.tool_call_id = tc.id @@ -726,39 +183,6 @@ LEFT JOIN LATERAL ( WHERE block.agent_id = result.agent_id AND block.owner_tool_call_result_id = result.id ) result_blocks ON true -WHERE tc.project_id = sqlc.arg(project_id) - AND tc.agent_id = sqlc.arg(agent_id) - AND tc.state = 'completed' - AND tc.source_event_sequence > sqlc.arg(after_event_sequence) - AND tool_result_event.sequence > sqlc.arg(after_event_sequence) - AND tc.source_event_sequence <= sqlc.arg(max_event_sequence) +WHERE tool_result_event.sequence > sqlc.arg(after_event_sequence) AND tool_result_event.sequence <= sqlc.arg(max_event_sequence) ORDER BY tc.source_event_sequence, call_block.ordinal, tc.created_at, tc.id; - --- name: CancelNonTerminalToolCallsForAgent :many -WITH locked_agent AS MATERIALIZED ( - SELECT agent.project_id, agent.id - FROM agents agent - WHERE agent.project_id = sqlc.arg(project_id) - AND agent.id = sqlc.arg(agent_id) - FOR UPDATE -) -UPDATE tool_calls call -SET state = 'completed', - runtime_lock_id = NULL -FROM locked_agent agent -CROSS JOIN tool_call_read_projection projection -WHERE call.agent_id = agent.id - AND call.state <> 'completed' - AND projection.project_id = agent.project_id - AND projection.agent_id = call.agent_id - AND projection.id = call.id - AND projection.turn_id = sqlc.arg(turn_id) -RETURNING call.id, projection.project_id, call.agent_id, - projection.turn_id, projection.source_event_id, projection.model_call_context_id, - call.provider_call_id, - call.name, call.input, - call.type, call.state, - 'canceled'::text AS outcome, call.runtime_lock_id, - '[]'::jsonb AS result_content_parts, - call.created_at; diff --git a/internal/storage/resource_limits_integration_test.go b/internal/storage/resource_limits_integration_test.go index 88b7c7113..cbf32e568 100644 --- a/internal/storage/resource_limits_integration_test.go +++ b/internal/storage/resource_limits_integration_test.go @@ -589,13 +589,14 @@ func TestAgentResourceLimitsPreserveReplays(t *testing.T) { t.Fatalf("create first agent: %v", err) } if _, err := pool.Exec(ctx, ` -INSERT INTO agents( - id, org_id, project_id, state, current_config_id, +WITH inserted AS (INSERT INTO agents( + id, root_agent_id, org_id, project_id, state, current_config_id, idempotency_key, created_at, updated_at ) -SELECT md5('resource-limit-agent-' || n)::uuid, $1, $2, 'active', $3, +SELECT md5('resource-limit-agent-' || n)::uuid, md5('resource-limit-agent-' || n)::uuid, $1, $2, 'active', $3, 'limit-seed-agent-' || n, $4, $4 -FROM generate_series(1, $5::integer) AS n +FROM generate_series(1, $5::integer) AS n RETURNING id) +INSERT INTO agent_execution_state(agent_id,turn_continuable,incomplete_tools) SELECT id,false,false FROM inserted `, testOrgID, testProjectID, configID, now, agentLimit-1); err != nil { t.Fatalf("seed agents to limit: %v", err) } @@ -773,11 +774,13 @@ func TestResourceLimitLocksDoNotBlockUnrelatedCreates(t *testing.T) { tokenCtx, cancelToken := context.WithTimeout(ctx, 5*time.Second) defer cancelToken() - if _, err := store.Execution().CreateBYOMachineDaemonToken(tokenCtx, executionstore.CreateBYOMachineDaemonTokenInput{ - OrgID: testOrgID, - MachineID: machine.ID, - Name: "Resource limit lock token", - }); err != nil { + if _, + err := store.Execution().CreateBYOMachineDaemonToken(tokenCtx, + executionstore.CreateBYOMachineDaemonTokenInput{ + OrgID: testOrgID, + MachineID: machine.ID, + Name: "Resource limit lock token", + }); err != nil { _ = blocker.Rollback(ctx) <-skillResult t.Fatalf("create daemon token while skill creation is blocked: %v", err) @@ -843,9 +846,10 @@ FROM generate_series(1, $3::integer) AS n t.Fatalf("machine over limit error = %v, want ErrConflict", err) } - firstToken, err := store.Execution().CreateBYOMachineDaemonToken(ctx, executionstore.CreateBYOMachineDaemonTokenInput{ - OrgID: testOrgID, MachineID: machine.ID, Name: "First daemon token", - }) + firstToken, err := store.Execution(). + CreateBYOMachineDaemonToken(ctx, executionstore.CreateBYOMachineDaemonTokenInput{ + OrgID: testOrgID, MachineID: machine.ID, Name: "First daemon token", + }) if err != nil { t.Fatalf("create first daemon token: %v", err) } diff --git a/internal/storage/secret_payload_current_integration_test.go b/internal/storage/secret_payload_current_integration_test.go new file mode 100644 index 000000000..35479fa6b --- /dev/null +++ b/internal/storage/secret_payload_current_integration_test.go @@ -0,0 +1,71 @@ +//go:build integration + +package storage + +import ( + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/secrets" + "github.com/omnara-ai/omnara/internal/storage/management" + "github.com/omnara-ai/omnara/internal/storage/secretstore" + "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/stretchr/testify/require" +) + +func TestCurrentOrgSecretPayloadTracksRotationAndChecksAuthority(t *testing.T) { + ctx := t.Context() + pool := openIntegrationDB(t, ctx) + seedMigratedDB(t, ctx, pool) + store := newSecretIntegrationStore(pool) + admin := createSecretTestUser(t, ctx, store, "Payload owner", "admin") + record, _, err := store.Secrets().CreateSecret(ctx, secretstore.CreateSecretInput{ + OrgID: testOrgID, OwnerKind: secretstore.SecretOwnerOrg, Name: "current-payload", + Material: secrets.GenericMaterial{Value: "initial"}, Actor: userPrincipal(admin.ID), + }) + require.NoError(t, err) + input := secretstore.ReadOrgOwnedSecretPayloadInput{ + OrgID: testOrgID, SecretID: record.ID, ManagementKind: management.Tenant, Kind: secretstore.SecretKindGeneric, + } + before := pool.Stat().AcquireCount() + payload, err := store.Secrets().ReadOrgOwnedSecretPayload(ctx, input) + require.NoError(t, err) + require.EqualValues(t, 1, pool.Stat().AcquireCount()-before) + require.Equal(t, "initial", payload.Payload[secrets.KeyValue]) + require.Equal(t, record.CurrentVersionID, payload.CurrentVersionID) + _, version, err := store.Secrets().CreateSecretVersion(ctx, secretstore.CreateSecretVersionInput{ + OrgID: testOrgID, SecretID: record.ID, + Material: secrets.GenericMaterial{Value: "rotated"}, Actor: userPrincipal(admin.ID), + }) + require.NoError(t, err) + payload, err = store.Secrets().ReadOrgOwnedSecretPayload(ctx, input) + require.NoError(t, err) + require.Equal(t, version.ID, payload.CurrentVersionID) + require.Equal(t, "rotated", payload.Payload[secrets.KeyValue]) + for _, test := range []struct { + name string + mutate func(*secretstore.ReadOrgOwnedSecretPayloadInput) + want error + }{ + {"organization", func(i *secretstore.ReadOrgOwnedSecretPayloadInput) { i.OrgID = uuid.New() }, storeerr.ErrNotFound}, + {"management", func(i *secretstore.ReadOrgOwnedSecretPayloadInput) { + i.ManagementKind = management.Cluster + }, storeerr.ErrNotFound}, + {"kind", func(i *secretstore.ReadOrgOwnedSecretPayloadInput) { + i.Kind = secretstore.SecretKindOAuthTokenSet + }, storeerr.ErrInvalidSecretRequest}, + } { + t.Run(test.name, func(t *testing.T) { + invalid := input + test.mutate(&invalid) + _, err := store.Secrets().ReadOrgOwnedSecretPayload(ctx, invalid) + require.ErrorIs(t, err, test.want) + }) + } + _, err = store.Secrets().DeleteSecret(ctx, secretstore.DeleteSecretInput{ + OrgID: testOrgID, SecretID: record.ID, Actor: userPrincipal(admin.ID), + }) + require.NoError(t, err) + _, err = store.Secrets().ReadOrgOwnedSecretPayload(ctx, input) + require.ErrorIs(t, err, storeerr.ErrNotFound) +} diff --git a/internal/storage/secretstore/secret_payload_store.go b/internal/storage/secretstore/secret_payload_store.go index 7e9a86475..cf0163c81 100644 --- a/internal/storage/secretstore/secret_payload_store.go +++ b/internal/storage/secretstore/secret_payload_store.go @@ -212,31 +212,24 @@ func (s *Store) ReadOrgOwnedSecretPayload( if err := management.Validate(input.ManagementKind); err != nil { return SecretPayloadRecord{}, invalidSecretRequest("%v", err) } - record, err := s.GetSecret(ctx, input.OrgID, input.SecretID) + row, err := s.q.GetCurrentSecretPayload(ctx, dbsqlc.GetCurrentSecretPayloadParams{ + OrgID: input.OrgID, ID: input.SecretID, + }) if err != nil { - return SecretPayloadRecord{}, err + if errors.Is(err, pgx.ErrNoRows) { + return SecretPayloadRecord{}, storeerr.ErrNotFound + } + return SecretPayloadRecord{}, fmt.Errorf("get current secret payload: %w", err) } - if record.OwnerKind != SecretOwnerOrg || record.ManagementKind != input.ManagementKind { + if row.OwnerKind != string(SecretOwnerOrg) || row.ManagementKind != string(input.ManagementKind) { return SecretPayloadRecord{}, storeerr.ErrNotFound } - if record.Kind != input.Kind { + if row.Kind != string(input.Kind) { return SecretPayloadRecord{}, invalidSecretRequest( - "secret kind %q does not match expected kind %q", - record.Kind, - input.Kind, + "secret kind %q does not match expected kind %q", row.Kind, input.Kind, ) } - version, err := s.q.GetSecretVersion( - ctx, - dbsqlc.GetSecretVersionParams{OrgID: input.OrgID, SecretID: input.SecretID, ID: record.CurrentVersionID}, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return SecretPayloadRecord{}, storeerr.ErrNotFound - } - return SecretPayloadRecord{}, fmt.Errorf("get current secret version: %w", err) - } - versionRecord := secretVersionFromGetSQLC(version) + versionRecord := secretVersionFromCurrentPayloadSQLC(row) payload, err := secrets.DecryptPayload( ctx, s.secretKeyWrapper, @@ -246,7 +239,7 @@ func (s *Store) ReadOrgOwnedSecretPayload( SecretID: versionRecord.SecretID.String(), VersionID: versionRecord.ID.String(), VersionNumber: versionRecord.VersionNumber, - Kind: record.Kind, + Kind: secrets.Kind(row.Kind), }, ) if err != nil { @@ -254,8 +247,8 @@ func (s *Store) ReadOrgOwnedSecretPayload( } return SecretPayloadRecord{ Payload: payload, - CurrentVersionID: record.CurrentVersionID, - OAuthAccessTokenExpires: version.OauthAccessTokenExpires, - OAuthAccessTokenRemaining: time.Duration(version.OauthAccessTokenRemainingSeconds * float64(time.Second)), + CurrentVersionID: row.ID, + OAuthAccessTokenExpires: row.OauthAccessTokenExpires, + OAuthAccessTokenRemaining: time.Duration(row.OauthAccessTokenRemainingSeconds * float64(time.Second)), }, nil } diff --git a/internal/storage/secretstore/secret_rows.go b/internal/storage/secretstore/secret_rows.go index 9c27b80d5..ab246f5a5 100644 --- a/internal/storage/secretstore/secret_rows.go +++ b/internal/storage/secretstore/secret_rows.go @@ -204,3 +204,21 @@ func secretGrantFromSQLC(row dbsqlc.SecretGrant) SecretGrantRecord { CreatedAt: row.CreatedAt, } } + +func secretVersionFromCurrentPayloadSQLC(row dbsqlc.GetCurrentSecretPayloadRow) SecretVersionRecord { + return SecretVersionRecord{ + ID: row.ID, + OrgID: row.OrgID, + SecretID: row.SecretID, + VersionNumber: row.VersionNumber, + PayloadKeys: append([]string(nil), row.PayloadKeys...), + EncryptionScheme: row.EncryptionScheme, + KeyID: row.KeyID, + DEKWrappedBy: row.DekWrappedBy, + EncryptedDEK: append([]byte(nil), row.EncryptedDek...), + EncryptedDEKNonce: append([]byte(nil), row.EncryptedDekNonce...), + Nonce: append([]byte(nil), row.Nonce...), + Ciphertext: append([]byte(nil), row.Ciphertext...), + CreatedAt: row.CreatedAt, + } +} diff --git a/internal/storage/store.go b/internal/storage/store.go index 2951e0ccb..225114443 100644 --- a/internal/storage/store.go +++ b/internal/storage/store.go @@ -83,12 +83,6 @@ func WithMachinePoolProviders(machinePoolProviders executionstore.MachinePoolPro } } -func WithModelCallRetryBackoff(backoff func(int, string) time.Duration) Option { - return func(config *storeConfig) { - config.modelCallRetryBackoff = backoff - } -} - func NewStore(pool *pgxpool.Pool, opts ...Option) *Store { config := storeConfig{} for _, opt := range opts { @@ -104,7 +98,9 @@ func NewStore(pool *pgxpool.Pool, opts ...Option) *Store { store.skills = skillstore.New(pool, config.blobs, store.identity) store.memories = memorystore.New(pool, config.memoryFS, config.memoryRecorder) store.artifacts = artifactstore.New(pool, config.blobs) - store.integrations = integrationstore.New(pool, executionstore.IntegrationAccess{}) + store.integrations = integrationstore.New(pool, integrationstore.Config{ + PostCommitPublisher: config.postCommitPublisher, Access: executionstore.IntegrationAccess{}, + }) store.execution = executionstore.New(pool, executionstore.Config{ PostCommitPublisher: config.postCommitPublisher, ModelCallRetryBackoff: config.modelCallRetryBackoff, diff --git a/internal/storage/store_integration.go b/internal/storage/store_integration.go new file mode 100644 index 000000000..f7d3d95df --- /dev/null +++ b/internal/storage/store_integration.go @@ -0,0 +1,11 @@ +//go:build integration + +package storage + +import "time" + +func WithModelCallRetryBackoff(backoff func(int, string) time.Duration) Option { + return func(config *storeConfig) { + config.modelCallRetryBackoff = backoff + } +} diff --git a/internal/storage/storeerr/errors.go b/internal/storage/storeerr/errors.go index 04b412f19..069b979c7 100644 --- a/internal/storage/storeerr/errors.go +++ b/internal/storage/storeerr/errors.go @@ -46,16 +46,18 @@ var ( ErrMachineSleepPendingWork = errors.New("machine has pending daemon work") ErrMachineNotWakeCapable = errors.New("machine is not wake capable") ErrMachineWakeUnresolved = errors.New("machine wake outcome is unresolved") - ErrProcessActionReportBlocked = errors.New("process action report is blocked by earlier non-terminal action") - ErrStateTransitionConflict = errors.New("state transition conflict") - ErrMachineProviderUnavailable = errors.New("machine provider unavailable") - ErrPermanentEnvironment = errors.New("permanent environment resolution error") - ErrModelGrantUnavailable = errors.New("model grant unavailable") - ErrUnauthorized = errors.New("unauthorized") - ErrOAuthScopeExceedsGrant = errors.New("requested scope exceeds the granted scope") - ErrNotFound = errors.New("not found") - ErrMCPOAuthFlowConsumed = errors.New("mcp oauth flow already consumed") - ErrIntegrationOAuthFlowConsumed = errors.New("integration oauth flow already consumed") + ErrProcessActionReportBlocked = errors.New( + "process action report is blocked by earlier non-terminal action", + ) + ErrStateTransitionConflict = errors.New("state transition conflict") + ErrMachineProviderUnavailable = errors.New("machine provider unavailable") + ErrPermanentEnvironment = errors.New("permanent environment resolution error") + ErrModelGrantUnavailable = errors.New("model grant unavailable") + ErrUnauthorized = errors.New("unauthorized") + ErrOAuthScopeExceedsGrant = errors.New("requested scope exceeds the granted scope") + ErrNotFound = errors.New("not found") + ErrMCPOAuthFlowConsumed = errors.New("mcp oauth flow already consumed") + ErrIntegrationOAuthFlowConsumed = errors.New("integration oauth flow already consumed") ) type FileContentConflictError struct { @@ -109,3 +111,5 @@ func InvalidRequest(err error) error { func IsNotFound(err error) bool { return errors.Is(err, ErrNotFound) || errors.Is(err, pgx.ErrNoRows) } + +var ErrCheckpointBoundaryUnsafe = errors.New("context checkpoint boundary is unsafe") diff --git a/migrations/000052_agent_execution_state.sql b/migrations/000052_agent_execution_state.sql new file mode 100644 index 000000000..4d2b9ebdd --- /dev/null +++ b/migrations/000052_agent_execution_state.sql @@ -0,0 +1,417 @@ +-- +goose Up + +LOCK TABLE agents, agent_inputs, agent_events, agent_turns, model_call_contexts, + model_outputs, tool_calls, tool_call_results, context_checkpoints IN ACCESS EXCLUSIVE MODE; + +ALTER TABLE agents ADD COLUMN root_agent_id uuid; +ALTER TABLE model_call_contexts + ADD COLUMN turn_id uuid, + ADD COLUMN opening_input_ids uuid[], + ADD COLUMN opening_event_sequence bigint; +ALTER TABLE agent_inputs + ADD COLUMN opening_input_ids uuid[], + ADD COLUMN opening_event_sequence bigint; +ALTER TABLE context_checkpoints + ADD COLUMN opening_input_ids uuid[], + ADD COLUMN opening_event_sequence bigint; + +CREATE TABLE agent_execution_state ( + agent_id uuid PRIMARY KEY REFERENCES agents(id) ON DELETE CASCADE, + current_turn_id uuid, + stop_sequence bigint NOT NULL DEFAULT 0 CHECK (stop_sequence >= 0), + answered_through_sequence bigint NOT NULL DEFAULT 0 CHECK (answered_through_sequence >= stop_sequence), + max_normal_input_sequence bigint NOT NULL DEFAULT 0 CHECK (max_normal_input_sequence >= 0), + max_context_input_sequence bigint NOT NULL DEFAULT 0 CHECK (max_context_input_sequence >= max_normal_input_sequence), + normal_context_id uuid, + compaction_context_id uuid, + pending_tool_output_id uuid, + pending_output_limit_id uuid, + pending_config_input_id uuid, + pending_checkpoint_id uuid, + turn_continuable boolean NOT NULL, + incomplete_tools boolean NOT NULL, + logical_ready_at timestamptz, + FOREIGN KEY (agent_id, current_turn_id) REFERENCES agent_turns(agent_id, id), + FOREIGN KEY (agent_id, normal_context_id) REFERENCES model_call_contexts(agent_id, id), + FOREIGN KEY (agent_id, compaction_context_id) REFERENCES model_call_contexts(agent_id, id), + FOREIGN KEY (agent_id, pending_tool_output_id) REFERENCES model_outputs(agent_id, id), + FOREIGN KEY (agent_id, pending_output_limit_id) REFERENCES model_outputs(agent_id, id), + FOREIGN KEY (agent_id, pending_config_input_id) REFERENCES agent_inputs(agent_id, id), + FOREIGN KEY (agent_id, pending_checkpoint_id) REFERENCES context_checkpoints(agent_id, id), + CHECK (compaction_context_id IS NULL OR normal_context_id IS NOT NULL), + CHECK (pending_tool_output_id IS NULL OR pending_output_limit_id IS NULL), + CHECK (NOT incomplete_tools OR turn_continuable), + CHECK (current_turn_id IS NOT NULL OR + (normal_context_id IS NULL AND compaction_context_id IS NULL AND pending_tool_output_id IS NULL + AND pending_output_limit_id IS NULL AND pending_config_input_id IS NULL AND pending_checkpoint_id IS NULL)) +); + +CREATE INDEX tool_calls_output_idx ON tool_calls(agent_id, model_output_id); +CREATE INDEX tool_calls_output_incomplete_idx ON tool_calls(agent_id, model_output_id) WHERE state <> 'completed'; +CREATE INDEX tool_calls_output_runnable_idx ON tool_calls(agent_id, model_output_id) + WHERE state = 'awaiting_authorization' OR (state = 'ready' AND type IN ('built_in', 'mcp')); +CREATE INDEX agent_events_output_boundary_idx ON agent_events(agent_id, sequence DESC) WHERE event_kind = 'model_output'; +CREATE INDEX model_call_contexts_turn_kind_idx + ON model_call_contexts(agent_id, turn_id, operation_kind, input_event_sequence DESC, attempt_number DESC); + +WITH RECURSIVE roots AS ( + SELECT id, project_id, id AS root_id FROM agents WHERE parent_agent_id IS NULL + UNION ALL + SELECT child.id, child.project_id, roots.root_id + FROM roots JOIN agents child ON child.project_id = roots.project_id AND child.parent_agent_id = roots.id +) +UPDATE agents SET root_agent_id = roots.root_id FROM roots WHERE agents.id = roots.id; + +-- +goose StatementBegin +DO $$ +DECLARE invalid_agent uuid; +BEGIN + SELECT id INTO invalid_agent FROM agents WHERE root_agent_id IS NULL LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'invalid agent ancestry: %', invalid_agent; END IF; + SELECT e.agent_id INTO invalid_agent FROM model_outputs o + JOIN model_call_contexts c ON c.agent_id=o.agent_id AND c.id=o.model_call_context_id + JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id + WHERE e.sequence<=c.input_event_sequence LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'model output precedes captured watermark: %', invalid_agent; END IF; +END; +$$; +-- +goose StatementEnd + +ALTER TABLE model_call_contexts DISABLE TRIGGER model_call_contexts_transition_guard; +ALTER TABLE agent_inputs DISABLE TRIGGER agent_inputs_mutation_policy; +ALTER TABLE context_checkpoints DISABLE TRIGGER context_checkpoints_immutable; + +UPDATE model_call_contexts c +SET turn_id = agent_turn_id_at_event_sequence(c.agent_id, c.input_event_sequence); +UPDATE model_call_contexts c SET opening_input_ids = o.ids, opening_event_sequence = o.first_sequence +FROM ( + SELECT c.id, coalesce(array_agg(o.input_id ORDER BY o.event_sequence) + FILTER (WHERE o.input_id IS NOT NULL), '{}'::uuid[]) AS ids, min(o.event_sequence) AS first_sequence + FROM model_call_contexts c + LEFT JOIN LATERAL agent_model_call_opening_content_inputs(c.project_id,c.agent_id,c.turn_id,c.input_event_sequence) o + ON true GROUP BY c.id +) o WHERE c.id = o.id; + +UPDATE agent_inputs i SET opening_input_ids = o.ids, opening_event_sequence = o.first_sequence +FROM ( + SELECT i.id, coalesce(array_agg(o.input_id ORDER BY o.event_sequence) + FILTER (WHERE o.input_id IS NOT NULL), '{}'::uuid[]) AS ids, min(o.event_sequence) AS first_sequence + FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id + LEFT JOIN LATERAL agent_model_call_opening_content_inputs(i.project_id,i.agent_id,e.turn_id,e.sequence) o ON true + WHERE i.input_kind='config_change' AND i.state='resolved' GROUP BY i.id +) o WHERE i.id=o.id; + +UPDATE context_checkpoints c SET opening_input_ids=o.ids, opening_event_sequence=o.first_sequence +FROM ( + SELECT c.id, coalesce(array_agg(o.input_id ORDER BY o.event_sequence) + FILTER (WHERE o.input_id IS NOT NULL), '{}'::uuid[]) AS ids, min(o.event_sequence) AS first_sequence + FROM context_checkpoints c JOIN agents a ON a.id=c.agent_id + JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id + LEFT JOIN LATERAL agent_model_call_opening_content_inputs(a.project_id,c.agent_id,e.turn_id,e.sequence) o ON true + GROUP BY c.id +) o WHERE c.id=o.id; + +ALTER TABLE model_call_contexts ENABLE TRIGGER model_call_contexts_transition_guard; +ALTER TABLE agent_inputs ENABLE TRIGGER agent_inputs_mutation_policy; +ALTER TABLE context_checkpoints ENABLE TRIGGER context_checkpoints_immutable; + +-- +goose StatementBegin +CREATE FUNCTION execution_opening_is_valid(ids uuid[], first_sequence bigint) +RETURNS boolean LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ + SELECT coalesce(ids IS NOT NULL AND array_ndims(ids) IS NOT DISTINCT FROM CASE WHEN cardinality(ids)=0 THEN NULL ELSE 1 END + AND (cardinality(ids)=0 AND first_sequence IS NULL OR cardinality(ids)>0 AND first_sequence>0) + AND NOT EXISTS(SELECT 1 FROM unnest(ids) id WHERE id IS NULL OR id='00000000-0000-0000-0000-000000000000'::uuid) + AND (cardinality(ids)=0 OR array_lower(ids,1)=1) + AND cardinality(ids)=(SELECT count(DISTINCT id) FROM unnest(ids) id), false) +$$; +-- +goose StatementEnd + +ALTER TABLE agents ALTER COLUMN root_agent_id SET NOT NULL, + ADD CONSTRAINT agents_root_fk FOREIGN KEY (project_id,root_agent_id) REFERENCES agents(project_id,id), + ADD CONSTRAINT agents_root_shape CHECK ((parent_agent_id IS NULL)=(root_agent_id=id)); +CREATE INDEX agents_root_idx ON agents(project_id,root_agent_id,id); +ALTER TABLE model_call_contexts ALTER COLUMN turn_id SET NOT NULL, + ALTER COLUMN opening_input_ids SET NOT NULL, ALTER COLUMN opening_event_sequence SET NOT NULL, + ADD CONSTRAINT model_call_contexts_turn_fk FOREIGN KEY (agent_id,turn_id) REFERENCES agent_turns(agent_id,id), + ADD CONSTRAINT model_call_contexts_opening_shape CHECK ( + cardinality(opening_input_ids)>0 AND execution_opening_is_valid(opening_input_ids,opening_event_sequence) + AND opening_event_sequence<=input_event_sequence); +ALTER TABLE agent_inputs ADD CONSTRAINT agent_inputs_opening_shape CHECK ( + CASE WHEN input_kind='config_change' AND state='resolved' + THEN execution_opening_is_valid(opening_input_ids,opening_event_sequence) + ELSE opening_input_ids IS NULL AND opening_event_sequence IS NULL END); +ALTER TABLE context_checkpoints ALTER COLUMN opening_input_ids SET NOT NULL, + ADD CONSTRAINT context_checkpoints_opening_shape CHECK (execution_opening_is_valid(opening_input_ids,opening_event_sequence)); + +-- +goose StatementBegin +CREATE FUNCTION enforce_execution_lineage_identity() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF TG_TABLE_NAME='agents' THEN + IF OLD.root_agent_id IS DISTINCT FROM NEW.root_agent_id THEN + RAISE EXCEPTION 'agent root is immutable' USING ERRCODE='25006'; + END IF; + ELSE + IF OLD.turn_id IS DISTINCT FROM NEW.turn_id + OR OLD.opening_input_ids IS DISTINCT FROM NEW.opening_input_ids + OR OLD.opening_event_sequence IS DISTINCT FROM NEW.opening_event_sequence THEN + RAISE EXCEPTION 'context lineage is immutable' USING ERRCODE='25006'; + END IF; + END IF; + RETURN NEW; +END; +$$; +-- +goose StatementEnd +CREATE TRIGGER agents_root_immutable BEFORE UPDATE OF root_agent_id ON agents + FOR EACH ROW EXECUTE FUNCTION enforce_execution_lineage_identity(); +CREATE TRIGGER model_call_contexts_lineage_immutable + BEFORE UPDATE OF turn_id,opening_input_ids,opening_event_sequence ON model_call_contexts + FOR EACH ROW EXECUTE FUNCTION enforce_execution_lineage_identity(); + +INSERT INTO agent_execution_state ( + agent_id,current_turn_id,stop_sequence,answered_through_sequence,max_normal_input_sequence,max_context_input_sequence, + normal_context_id,compaction_context_id,turn_continuable,incomplete_tools,logical_ready_at) +SELECT a.id, t.id, coalesce(s.sequence,0), greatest(coalesce(s.sequence,0),coalesce(answer.sequence,0)), + coalesce(w.normal_sequence,0), coalesce(w.context_sequence,0), n.id,c.id, + a.state='active' AND (agent_has_incomplete_tool_batch(a.project_id,a.id) + OR EXISTS(SELECT 1 FROM agent_next_model_work(a.project_id,a.id)) + OR EXISTS(SELECT 1 FROM agent_continuable_model_contexts(a.project_id,a.id))), + a.state='active' AND agent_has_incomplete_tool_batch(a.project_id,a.id), + CASE WHEN a.state='active' THEN agent_next_wakeup_ready_at(a.project_id,a.id) END +FROM agents a +LEFT JOIN LATERAL (SELECT id FROM agent_turns WHERE agent_id=a.id ORDER BY turn_sequence DESC LIMIT 1) t ON true +LEFT JOIN LATERAL (SELECT max(sequence) AS sequence FROM agent_stop_events WHERE agent_id=a.id AND project_id=a.project_id) s ON true +LEFT JOIN LATERAL (SELECT sequence FROM agent_events WHERE agent_id=a.id AND event_kind='model_output' ORDER BY sequence DESC LIMIT 1) answer ON true +LEFT JOIN LATERAL (SELECT max(input_event_sequence) FILTER (WHERE operation_kind='normal') AS normal_sequence, + max(input_event_sequence) AS context_sequence FROM model_call_contexts WHERE agent_id=a.id) w ON true +LEFT JOIN LATERAL (SELECT id,input_event_sequence FROM model_call_contexts WHERE agent_id=a.id AND turn_id=t.id + AND operation_kind='normal' ORDER BY input_event_sequence DESC,attempt_number DESC LIMIT 1) n ON true +LEFT JOIN LATERAL (SELECT id FROM model_call_contexts WHERE agent_id=a.id AND turn_id=t.id + AND operation_kind='compaction' AND input_event_sequence=n.input_event_sequence + ORDER BY source_event_sequence_end,attempt_number DESC LIMIT 1) c ON true; + +CREATE TEMP TABLE execution_pending_batches ON COMMIT DROP AS +SELECT DISTINCT a.id AS agent_id, eligible.model_output_id +FROM agents a CROSS JOIN LATERAL ( + SELECT c.model_output_id FROM tool_calls c + JOIN agent_events e ON e.agent_id=c.agent_id AND e.model_output_id=c.model_output_id + JOIN agent_execution_state h ON h.agent_id=c.agent_id AND h.current_turn_id=e.turn_id + WHERE c.agent_id=a.id AND c.state<>'completed' AND e.sequence>=h.stop_sequence + UNION + SELECT f.model_output_id FROM agent_model_result_frontiers(a.project_id,a.id) f +) eligible; + +-- +goose StatementBegin +DO $$ +DECLARE invalid_agent uuid; +BEGIN + SELECT agent_id INTO invalid_agent FROM execution_pending_batches GROUP BY agent_id HAVING count(*)>1 LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'multiple unconsumed tool outputs: %', invalid_agent; END IF; + SELECT a.id INTO invalid_agent FROM agents a + WHERE (SELECT count(*) FROM agent_continuable_model_contexts(a.project_id,a.id))>1 LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'multiple governing contexts: %', invalid_agent; END IF; + SELECT a.id INTO invalid_agent FROM agents a JOIN agent_execution_state h ON h.agent_id=a.id + WHERE EXISTS(SELECT 1 FROM agent_continuable_model_contexts(a.project_id,a.id) c + WHERE c.model_call_context_id IS DISTINCT FROM h.normal_context_id + AND c.model_call_context_id IS DISTINCT FROM h.compaction_context_id) LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'unrepresented governing context: %', invalid_agent; END IF; +END; +$$; +-- +goose StatementEnd + +UPDATE agent_execution_state h SET pending_tool_output_id=b.model_output_id +FROM execution_pending_batches b WHERE b.agent_id=h.agent_id; +UPDATE agent_execution_state h SET pending_config_input_id=e.agent_input_id +FROM agents a CROSS JOIN LATERAL agent_unconsumed_config_change_frontiers(a.project_id,a.id) f +JOIN agent_events e ON e.agent_id=a.id AND e.sequence=f.config_event_sequence +WHERE h.agent_id=a.id; +UPDATE agent_execution_state h SET pending_checkpoint_id=e.context_checkpoint_id +FROM agents a CROSS JOIN LATERAL agent_unconsumed_context_checkpoint_frontiers(a.project_id,a.id) f +JOIN agent_events e ON e.agent_id=a.id AND e.sequence=f.checkpoint_event_sequence +WHERE h.agent_id=a.id; + +CREATE TEMP TABLE execution_pending_limits ON COMMIT DROP AS +SELECT h.agent_id,o.id FROM agent_execution_state h +JOIN agent_events e ON e.agent_id=h.agent_id AND e.turn_id=h.current_turn_id AND e.event_kind='model_output' +JOIN model_outputs o ON o.agent_id=e.agent_id AND o.id=e.model_output_id +JOIN model_call_contexts c ON c.agent_id=o.agent_id AND c.id=o.model_call_context_id +WHERE o.stop_reason='max_tokens' AND c.operation_kind='normal' AND c.state='succeeded' + AND e.sequence>h.max_normal_input_sequence AND e.sequence>=h.stop_sequence + AND NOT EXISTS(SELECT 1 FROM tool_calls tc WHERE tc.agent_id=o.agent_id AND tc.model_output_id=o.id); + +-- +goose StatementBegin +DO $$ +DECLARE invalid_agent uuid; +BEGIN + SELECT agent_id INTO invalid_agent FROM execution_pending_limits GROUP BY agent_id HAVING count(*)>1 LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'multiple unconsumed output limits: %', invalid_agent; END IF; +END; +$$; +-- +goose StatementEnd +UPDATE agent_execution_state h SET pending_output_limit_id=o.id FROM execution_pending_limits o WHERE o.agent_id=h.agent_id; + +-- +goose StatementBegin +CREATE FUNCTION enforce_execution_root() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF NEW.parent_agent_id IS NOT NULL AND NOT EXISTS ( + SELECT 1 FROM agents parent WHERE parent.project_id=NEW.project_id AND parent.id=NEW.parent_agent_id + AND parent.root_agent_id=NEW.root_agent_id + ) THEN + RAISE EXCEPTION 'agent % root must match parent', NEW.id USING ERRCODE='23514'; + END IF; + RETURN NEW; +END; +$$; +-- +goose StatementEnd +CREATE TRIGGER agents_root_valid BEFORE INSERT ON agents FOR EACH ROW EXECUTE FUNCTION enforce_execution_root(); + +-- +goose StatementBegin +CREATE FUNCTION enforce_execution_context_turn() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF NEW.turn_id IS DISTINCT FROM ( + SELECT turn_id FROM agent_events WHERE agent_id=NEW.agent_id AND is_opening_event + AND sequence<=NEW.input_event_sequence ORDER BY sequence DESC LIMIT 1 + ) THEN + RAISE EXCEPTION 'context % has invalid captured turn', NEW.id USING ERRCODE='23514'; + END IF; + RETURN NEW; +END; +$$; +-- +goose StatementEnd +CREATE TRIGGER model_call_contexts_execution_turn_valid BEFORE INSERT ON model_call_contexts + FOR EACH ROW EXECUTE FUNCTION enforce_execution_context_turn(); + +DELETE FROM agent_wakeups w USING agents a,agent_execution_state h +WHERE w.agent_id=a.id AND h.agent_id=a.id AND (a.state='archived' OR h.logical_ready_at IS NULL + OR EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id)); +INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) +SELECT h.agent_id,h.logical_ready_at,statement_timestamp() +FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id +WHERE a.state='active' AND h.logical_ready_at IS NOT NULL + AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id) +ON CONFLICT(agent_id) DO UPDATE SET ready_at=CASE + WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() + THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END, + updated_at=excluded.updated_at; + +-- +goose StatementBegin +DO $$ +DECLARE invalid_agent uuid; +BEGIN + SELECT a.id INTO invalid_agent FROM agents a LEFT JOIN agent_execution_state h ON h.agent_id=a.id + WHERE h.agent_id IS NULL LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'missing execution state: %', invalid_agent; END IF; + SELECT h.agent_id INTO invalid_agent FROM agent_execution_state h JOIN agents a ON a.id=h.agent_id + WHERE a.state='active' AND ( + h.incomplete_tools IS DISTINCT FROM agent_has_incomplete_tool_batch(a.project_id,a.id) + OR h.turn_continuable IS DISTINCT FROM (agent_has_incomplete_tool_batch(a.project_id,a.id) + OR EXISTS(SELECT 1 FROM agent_next_model_work(a.project_id,a.id)) + OR EXISTS(SELECT 1 FROM agent_continuable_model_contexts(a.project_id,a.id))) + OR (h.logical_ready_at IS NULL) IS DISTINCT FROM (agent_next_wakeup_ready_at(a.project_id,a.id) IS NULL) + OR ((h.logical_ready_at>statement_timestamp() OR agent_next_wakeup_ready_at(a.project_id,a.id)>statement_timestamp()) + AND h.logical_ready_at IS DISTINCT FROM agent_next_wakeup_ready_at(a.project_id,a.id)) + ) LIMIT 1; + IF FOUND THEN RAISE EXCEPTION 'execution summary mismatch: %', invalid_agent; END IF; +END; +$$; +-- +goose StatementEnd + +-- +goose StatementBegin +CREATE OR REPLACE FUNCTION enforce_agent_event_payload_turn() +RETURNS trigger +LANGUAGE plpgsql +AS $$ +DECLARE + expected_turn_id uuid; +BEGIN + IF NEW.event_kind = 'model_output' THEN + SELECT context_turn.turn_id + INTO expected_turn_id + FROM model_outputs output + JOIN model_call_contexts context_turn ON context_turn.agent_id = output.agent_id + AND context_turn.id = output.model_call_context_id + WHERE output.agent_id = NEW.agent_id + AND output.id = NEW.model_output_id; + + IF expected_turn_id IS NULL OR expected_turn_id IS DISTINCT FROM NEW.turn_id THEN + RAISE EXCEPTION 'model_output event % must use turn % captured by model context', NEW.id, expected_turn_id + USING ERRCODE = '23514'; + END IF; + ELSIF NEW.event_kind = 'tool_result' THEN + SELECT source_event.turn_id + INTO expected_turn_id + FROM tool_call_results result + JOIN tool_calls tool_call ON tool_call.agent_id = result.agent_id + AND tool_call.id = result.tool_call_id + JOIN agent_events source_event ON source_event.agent_id = tool_call.agent_id + AND source_event.model_output_id = tool_call.model_output_id + AND source_event.event_kind = 'model_output' + WHERE result.agent_id = NEW.agent_id + AND result.id = NEW.tool_call_result_id; + + IF expected_turn_id IS NULL OR expected_turn_id IS DISTINCT FROM NEW.turn_id THEN + RAISE EXCEPTION 'tool_result event % must use turn % derived from tool call source event', NEW.id, expected_turn_id + USING ERRCODE = '23514'; + END IF; + ELSIF NEW.event_kind = 'context_checkpoint' THEN + SELECT context_turn.turn_id + INTO expected_turn_id + FROM context_checkpoints checkpoint + JOIN model_call_contexts context_turn ON context_turn.agent_id = checkpoint.agent_id + AND context_turn.id = checkpoint.producer_model_call_context_id + WHERE checkpoint.agent_id = NEW.agent_id + AND checkpoint.id = NEW.context_checkpoint_id; + + IF expected_turn_id IS NULL OR expected_turn_id IS DISTINCT FROM NEW.turn_id THEN + RAISE EXCEPTION 'context_checkpoint event % must use turn % derived from producer context', NEW.id, expected_turn_id + USING ERRCODE = '23514'; + END IF; + END IF; + + RETURN NEW; +END; +$$; +-- +goose StatementEnd + +CREATE OR REPLACE VIEW tool_call_read_projection AS +SELECT tool_call.id, + context.project_id, + tool_call.agent_id, + context.turn_id, + source_event.id AS source_event_id, + source_event.sequence AS source_event_sequence, + model_output.model_call_context_id, + tool_call.model_output_id, + tool_call.provider_call_id, + tool_call.name, + tool_call.input, + tool_call.type, + tool_call.state, + tool_call.runtime_lock_id, + tool_call.created_at +FROM tool_calls tool_call +JOIN model_outputs model_output ON model_output.agent_id = tool_call.agent_id + AND model_output.id = tool_call.model_output_id +JOIN model_call_contexts context ON context.agent_id = model_output.agent_id + AND context.id = model_output.model_call_context_id +JOIN agent_events source_event ON source_event.agent_id = tool_call.agent_id + AND source_event.model_output_id = tool_call.model_output_id + AND source_event.event_kind = 'model_output'; + +DROP VIEW model_call_context_turns; +DROP VIEW agent_stop_events; +DROP FUNCTION agent_next_wakeup_ready_at(uuid,uuid); +DROP FUNCTION agent_next_model_work(uuid,uuid); +DROP FUNCTION agent_model_result_frontiers(uuid,uuid); +DROP FUNCTION agent_has_incomplete_tool_batch(uuid,uuid); +DROP FUNCTION agent_tool_work_frontiers(uuid,uuid); +DROP FUNCTION agent_unconsumed_config_change_frontiers(uuid,uuid); +DROP FUNCTION agent_unconsumed_context_checkpoint_frontiers(uuid,uuid); +DROP FUNCTION agent_continuable_model_contexts(uuid,uuid); +DROP FUNCTION model_call_context_has_later_semantic_event(uuid,uuid,uuid); +DROP FUNCTION agent_latest_unstarted_model_call_opening_inputs(uuid,uuid); +DROP FUNCTION agent_model_call_opening_content_inputs(uuid,uuid,uuid,bigint); +DROP FUNCTION agent_latest_turn_id(uuid,uuid); +DROP FUNCTION agent_turn_opening_content_inputs(uuid,uuid,bigint); +DROP FUNCTION agent_turn_id_at_event_sequence(uuid,bigint); + +ALTER TABLE agent_wakeups DROP COLUMN metadata; diff --git a/migrations/000053_agent_payload_query_plans.sql b/migrations/000053_agent_payload_query_plans.sql new file mode 100644 index 000000000..e1ce24670 --- /dev/null +++ b/migrations/000053_agent_payload_query_plans.sql @@ -0,0 +1,30 @@ +-- +goose Up + +-- +goose StatementBegin +CREATE FUNCTION agent_payload_created_at( + p_project_id uuid, p_agent_id uuid, p_kind text, p_id uuid, p_admission boolean +) +RETURNS timestamptz +LANGUAGE plpgsql STABLE +AS $$ +DECLARE result timestamptz; +BEGIN + IF NOT EXISTS (SELECT 1 FROM agents WHERE project_id = p_project_id AND id = p_agent_id) THEN + RETURN NULL; + END IF; + CASE p_kind + WHEN 'agent_input' THEN + SELECT CASE WHEN p_admission THEN statement_timestamp() ELSE queued_at END INTO result + FROM agent_inputs WHERE agent_id = p_agent_id AND id = p_id; + WHEN 'model_output' THEN + SELECT created_at INTO result FROM model_outputs WHERE agent_id = p_agent_id AND id = p_id; + WHEN 'tool_call_result' THEN + SELECT completed_at INTO result FROM tool_call_results WHERE agent_id = p_agent_id AND id = p_id; + WHEN 'context_checkpoint' THEN + SELECT created_at INTO result FROM context_checkpoints WHERE agent_id = p_agent_id AND id = p_id; + ELSE RETURN NULL; + END CASE; + RETURN result; +END; +$$; +-- +goose StatementEnd diff --git a/sqlc.vet-db.yaml b/sqlc.vet-db.yaml index ff4e287e3..e6e52ce22 100644 --- a/sqlc.vet-db.yaml +++ b/sqlc.vet-db.yaml @@ -62,7 +62,7 @@ sql: queries: "internal/storage/queries" database: uri: "${SQLC_DATABASE_URL}" - rules: + rules: &postgres_rules - sqlc/db-prepare - users-deleted-at - orgs-deleted-at @@ -90,7 +90,31 @@ sql: query_parameter_limit: 0 output_db_file_name: "db_gen.go" output_models_file_name: "models_gen.go" - overrides: + overrides: &postgres_overrides + - column: "processes.execution_spec" + go_type: + import: "github.com/omnara-ai/omnara/internal/processcmd" + type: "ExecutionSpec" + - db_type: "pg_catalog.uuid" + go_type: + import: "github.com/google/uuid" + type: "UUID" + - db_type: "pg_catalog.uuid" + nullable: true + go_type: + import: "github.com/google/uuid" + type: "UUID" + pointer: true + - db_type: "uuid" + go_type: + import: "github.com/google/uuid" + type: "UUID" + - db_type: "uuid" + nullable: true + go_type: + import: "github.com/google/uuid" + type: "UUID" + pointer: true - db_type: "pg_catalog.timestamptz" go_type: import: "time" @@ -166,3 +190,23 @@ sql: import: "encoding/json" type: "RawMessage" pointer: true + - engine: "postgresql" + schema: "migrations" + queries: "internal/storage/internal/agentexecution/queries" + database: + uri: "${SQLC_DATABASE_URL}" + rules: *postgres_rules + gen: + go: + package: "executiondb" + emit_methods_with_db_argument: true + out: "internal/storage/internal/agentexecution/internal/executiondb" + sql_package: "pgx/v5" + emit_empty_slices: true + emit_json_tags: false + emit_interface: false + omit_unused_structs: true + query_parameter_limit: 0 + output_db_file_name: "db_gen.go" + output_models_file_name: "models_gen.go" + overrides: *postgres_overrides diff --git a/sqlc.yaml b/sqlc.yaml index 3073b3992..ad3520e8f 100644 --- a/sqlc.yaml +++ b/sqlc.yaml @@ -64,7 +64,7 @@ sql: - engine: "postgresql" schema: "migrations" queries: "internal/storage/queries" - rules: + rules: &postgres_rules - users-deleted-at - orgs-deleted-at - projects-deleted-at @@ -91,7 +91,7 @@ sql: query_parameter_limit: 0 output_db_file_name: "db_gen.go" output_models_file_name: "models_gen.go" - overrides: + overrides: &postgres_overrides - column: "processes.execution_spec" go_type: import: "github.com/omnara-ai/omnara/internal/processcmd" @@ -191,6 +191,24 @@ sql: import: "encoding/json" type: "RawMessage" pointer: true + - engine: "postgresql" + schema: "migrations" + queries: "internal/storage/internal/agentexecution/queries" + rules: *postgres_rules + gen: + go: + package: "executiondb" + emit_methods_with_db_argument: true + out: "internal/storage/internal/agentexecution/internal/executiondb" + sql_package: "pgx/v5" + emit_empty_slices: true + emit_json_tags: false + emit_interface: false + omit_unused_structs: true + query_parameter_limit: 0 + output_db_file_name: "db_gen.go" + output_models_file_name: "models_gen.go" + overrides: *postgres_overrides - engine: "sqlite" schema: "internal/machinedaemon/statedb/migrations" queries: "internal/machinedaemon/statedb/queries" diff --git a/tools/ci/sqlrules/checker.go b/tools/ci/sqlrules/checker.go index 430081162..ee39f704a 100644 --- a/tools/ci/sqlrules/checker.go +++ b/tools/ci/sqlrules/checker.go @@ -235,7 +235,9 @@ func reconcileQueryMetadata(body []byte, start int, metadata queryMetadata) (que sqlcMetadata, err := parseSQLCQueryMetadata(body, start) if err != nil { metadata.diagnostic = sqlcMetadata.diagnostic - return metadata, &queryOwnershipError{message: fmt.Sprintf("sqlc rejects query ownership metadata: %v", err)} + return metadata, &queryOwnershipError{ + message: fmt.Sprintf("sqlc rejects query ownership metadata: %v", err), + } } if metadata.name == sqlcMetadata.name && metadata.command == sqlcMetadata.command { return metadata, nil @@ -431,6 +433,7 @@ func (checker *queryChecker) checkMessage( ) } case *pg_query.InsertStmt: + checker.checkExecutionOwnership(node.GetRelation()) if len(node.GetCols()) == 0 && node.GetSelectStmt() != nil { checker.report( ruleExplicitInsertColumns, @@ -444,6 +447,7 @@ func (checker *queryChecker) checkMessage( } checker.checkOutputTargets(node.GetReturningList(), mutationRelationNames(node.GetRelation(), nil)) case *pg_query.UpdateStmt: + checker.checkExecutionOwnership(node.GetRelation()) if node.GetWhereClause() == nil { checker.report( ruleMutationPredicate, @@ -458,6 +462,7 @@ func (checker *queryChecker) checkMessage( mutationRelationNames(node.GetRelation(), node.GetFromClause()), ) case *pg_query.DeleteStmt: + checker.checkExecutionOwnership(node.GetRelation()) if node.GetWhereClause() == nil { checker.report( ruleMutationPredicate, @@ -544,7 +549,10 @@ func (checker *queryChecker) checkInsertDatabaseOwnedTime(statement *pg_query.In checker.checkInsertSelectionDatabaseOwnedTime(columns, selection) } -func (checker *queryChecker) checkInsertSelectionDatabaseOwnedTime(columns []string, selection *pg_query.SelectStmt) { +func (checker *queryChecker) checkInsertSelectionDatabaseOwnedTime( + columns []string, + selection *pg_query.SelectStmt, +) { if selection == nil { return } @@ -781,22 +789,26 @@ func projectionContainsImplicitShape(expression *pg_query.Node, relationNames ma return false } found := false - walkMessage(expression.ProtoReflect(), nil, func(candidate protoreflect.Message, ancestors []protoreflect.Message) { - if found || hasSelectAncestor(ancestors) { - return - } - switch node := candidate.Interface().(type) { - case *pg_query.A_Star: - found = true - case *pg_query.FuncCall: - found = isSQLCEmbedCall(node) - case *pg_query.ColumnRef: - path := nodePath(node.GetFields()) - if len(path) == 1 { - _, found = relationNames[path[0]] + walkMessage( + expression.ProtoReflect(), + nil, + func(candidate protoreflect.Message, ancestors []protoreflect.Message) { + if found || hasSelectAncestor(ancestors) { + return } - } - }) + switch node := candidate.Interface().(type) { + case *pg_query.A_Star: + found = true + case *pg_query.FuncCall: + found = isSQLCEmbedCall(node) + case *pg_query.ColumnRef: + path := nodePath(node.GetFields()) + if len(path) == 1 { + _, found = relationNames[path[0]] + } + } + }, + ) return found } @@ -961,11 +973,15 @@ func expressionContainsConstant(expression *pg_query.Node, constant *pg_query.A_ return false } found := false - walkMessage(expression.ProtoReflect(), nil, func(candidate protoreflect.Message, _ []protoreflect.Message) { - if node, ok := candidate.Interface().(*pg_query.A_Const); ok && node == constant { - found = true - } - }) + walkMessage( + expression.ProtoReflect(), + nil, + func(candidate protoreflect.Message, _ []protoreflect.Message) { + if node, ok := candidate.Interface().(*pg_query.A_Const); ok && node == constant { + found = true + } + }, + ) return found } @@ -1090,19 +1106,23 @@ func expressionContainsApplicationTime(expression *pg_query.Node) bool { } found := false - walkMessage(expression.ProtoReflect(), nil, func(candidate protoreflect.Message, _ []protoreflect.Message) { - if found { - return - } - switch node := candidate.Interface().(type) { - case *pg_query.FuncCall: - found = isSQLCParameterCall(node) - case *pg_query.ParamRef: - found = true - case *pg_query.A_Expr: - found = isAtParameter(node) - } - }) + walkMessage( + expression.ProtoReflect(), + nil, + func(candidate protoreflect.Message, _ []protoreflect.Message) { + if found { + return + } + switch node := candidate.Interface().(type) { + case *pg_query.FuncCall: + found = isSQLCParameterCall(node) + case *pg_query.ParamRef: + found = true + case *pg_query.A_Expr: + found = isAtParameter(node) + } + }, + ) return found } @@ -1125,7 +1145,8 @@ func isDurationArithmetic(expression *pg_query.A_Expr) bool { if len(operator) != 1 || (operator[0] != "*" && operator[0] != "/") { return false } - return expressionContainsInterval(expression.GetLexpr()) || expressionContainsInterval(expression.GetRexpr()) + return expressionContainsInterval(expression.GetLexpr()) || + expressionContainsInterval(expression.GetRexpr()) } func expressionContainsInterval(expression *pg_query.Node) bool { @@ -1133,17 +1154,21 @@ func expressionContainsInterval(expression *pg_query.Node) bool { return false } found := false - walkMessage(expression.ProtoReflect(), nil, func(candidate protoreflect.Message, _ []protoreflect.Message) { - if found { - return - } - switch node := candidate.Interface().(type) { - case *pg_query.TypeCast: - found = typeNameEndsWith(node.GetTypeName(), "interval") - case *pg_query.FuncCall: - found = isDurationConstructor(node) - } - }) + walkMessage( + expression.ProtoReflect(), + nil, + func(candidate protoreflect.Message, _ []protoreflect.Message) { + if found { + return + } + switch node := candidate.Interface().(type) { + case *pg_query.TypeCast: + found = typeNameEndsWith(node.GetTypeName(), "interval") + case *pg_query.FuncCall: + found = isDurationConstructor(node) + } + }, + ) return found } @@ -1210,3 +1235,30 @@ func sourcePosition(source []byte, offset int) (int, int) { lineStart := bytes.LastIndexByte(source[:offset], '\n') + 1 return line, offset - lineStart + 1 } + +func (checker *queryChecker) checkExecutionOwnership(relation *pg_query.RangeVar) { + if strings.Contains(filepath.ToSlash(checker.path), "internal/storage/internal/agentexecution/queries/") { + return + } + switch relation.GetRelname() { + case "agents", + "agent_inputs", + "agent_events", + "agent_turns", + "model_call_contexts", + "model_outputs", + "tool_calls", + "tool_call_results", + "context_checkpoints", + "content_blocks", + "agent_interactions", + "agent_execution_state", + "agent_wakeups", + "agent_runtime_locks": + checker.report( + "execution-dml-ownership", + "execution writes belong to agentexecution/queries", + relation.GetLocation(), + ) + } +} diff --git a/tools/ci/sqlrules/checker_test.go b/tools/ci/sqlrules/checker_test.go index fd8d93fc3..bfc9eaf9a 100644 --- a/tools/ci/sqlrules/checker_test.go +++ b/tools/ci/sqlrules/checker_test.go @@ -92,7 +92,11 @@ func TestCheckerRules(t *testing.T) { sql: "SELECT id FROM jobs WHERE created_at < 'now'", }, {name: "today label", command: "many", sql: "SELECT id FROM jobs WHERE label = 'today'"}, - {name: "tomorrow JSON value", command: "one", sql: "SELECT jsonb_build_object('schedule', 'tomorrow')"}, + { + name: "tomorrow JSON value", + command: "one", + sql: "SELECT jsonb_build_object('schedule', 'tomorrow')", + }, { name: "temporal conversion around text comparison", command: "many", @@ -219,18 +223,36 @@ func TestCheckerRules(t *testing.T) { sql: "WITH locked AS (SELECT id FROM jobs WHERE id = sqlc.arg(id) FOR UPDATE) " + "UPDATE jobs SET state = 'done' FROM locked WHERE jobs.id = locked.id", }, - reject("update without predicate", "exec", "UPDATE users SET display_name = 'updated'", ruleMutationPredicate), + reject( + "update without predicate", + "exec", + "UPDATE users SET display_name = 'updated'", + ruleMutationPredicate, + ), reject("delete without predicate", "exec", "DELETE FROM users", ruleMutationPredicate), - {name: "update with predicate", command: "exec", sql: "UPDATE users SET display_name = 'updated' WHERE id = 1"}, + { + name: "update with predicate", + command: "exec", + sql: "UPDATE users SET display_name = 'updated' WHERE id = 1", + }, {name: "delete with predicate", command: "exec", sql: "DELETE FROM users WHERE id = 1"}, reject("implicit insert columns", "exec", "INSERT INTO users VALUES (1)", ruleExplicitInsertColumns), - reject("implicit insert select columns", "exec", "INSERT INTO users SELECT 1", ruleExplicitInsertColumns), + reject( + "implicit insert select columns", + "exec", + "INSERT INTO users SELECT 1", + ruleExplicitInsertColumns, + ), {name: "default values", command: "exec", sql: "INSERT INTO users DEFAULT VALUES"}, {name: "explicit insert columns", command: "exec", sql: "INSERT INTO users (id) VALUES (1)"}, reject("comma join", "many", "SELECT first.id FROM first, second", ruleExplicitJoins), reject("natural join", "many", "SELECT first.id FROM first NATURAL JOIN second", ruleExplicitJoins), {name: "cross join", command: "many", sql: "SELECT first.id FROM first CROSS JOIN second"}, - {name: "join on", command: "many", sql: "SELECT first.id FROM first JOIN second ON second.id = first.id"}, + { + name: "join on", + command: "many", + sql: "SELECT first.id FROM first JOIN second ON second.id = first.id", + }, {name: "join using", command: "many", sql: "SELECT first.id FROM first JOIN second USING (id)"}, reject( "unaliased relation function", @@ -295,9 +317,22 @@ func TestCheckerRules(t *testing.T) { "SELECT value FROM (SELECT 1 AS value)", ruleExplicitDerivedAliases, ), - {name: "aliased derived subquery", command: "one", sql: "SELECT derived.value FROM (SELECT 1 AS value) AS derived"}, - reject("unordered limited collection", "many", "SELECT id FROM users LIMIT 10", ruleLimitedCollectionOrder), - {name: "ordered limited collection", command: "many", sql: "SELECT id FROM users ORDER BY id LIMIT 10"}, + { + name: "aliased derived subquery", + command: "one", + sql: "SELECT derived.value FROM (SELECT 1 AS value) AS derived", + }, + reject( + "unordered limited collection", + "many", + "SELECT id FROM users LIMIT 10", + ruleLimitedCollectionOrder, + ), + { + name: "ordered limited collection", + command: "many", + sql: "SELECT id FROM users ORDER BY id LIMIT 10", + }, reject( "unordered limited batch collection", "batchmany", @@ -315,7 +350,12 @@ func TestCheckerRules(t *testing.T) { command: "many", sql: "SELECT id FROM users WHERE EXISTS (SELECT 1 FROM projects LIMIT 1)", }, - reject("offset pagination", "many", "SELECT id FROM users ORDER BY id OFFSET 10", ruleNoOffsetPagination), + reject( + "offset pagination", + "many", + "SELECT id FROM users ORDER BY id OFFSET 10", + ruleNoOffsetPagination, + ), reject("utility statement", "exec", "TRUNCATE users", ruleProductDMLOnly), reject("select into", "exec", "SELECT id INTO archived_users FROM users", ruleProductDMLOnly), { @@ -410,7 +450,10 @@ func TestCheckerRejectsFileContainingOnlyUnnamedSQL(t *testing.T) { } func TestCheckerAllowsCommentOnlyPreamble(t *testing.T) { - issues := checkSource("fixture.sql", []byte("-- product query definitions\n\n-- name: Safe :one\nSELECT 1;\n")) + issues := checkSource( + "fixture.sql", + []byte("-- product query definitions\n\n-- name: Safe :one\nSELECT 1;\n"), + ) if len(issues) != 0 { t.Fatalf("issues = %+v, want none", issues) } @@ -424,9 +467,12 @@ func TestParseErrorDiagnostic(t *testing.T) { } func TestSecondStatementRequiresItsOwnQueryAnnotation(t *testing.T) { - source := []byte("-- name: Multiple :many\nSELECT id FROM users ORDER BY id LIMIT 1;\nSELECT id FROM users;\n") + source := []byte( + "-- name: Multiple :many\nSELECT id FROM users ORDER BY id LIMIT 1;\nSELECT id FROM users;\n", + ) issues := checkSource("fixture.sql", source) - if len(issues) != 1 || issues[0].Rule != ruleNamedQueryRequired || issues[0].Line != 3 || issues[0].Column != 1 { + if len(issues) != 1 || issues[0].Rule != ruleNamedQueryRequired || issues[0].Line != 3 || + issues[0].Column != 1 { t.Fatalf("issues = %+v, want a line 3:1 %s violation", issues, ruleNamedQueryRequired) } } @@ -436,7 +482,15 @@ func TestProductionQueryCorpus(t *testing.T) { if !ok { t.Fatal("resolve test file path") } - queryDirectory := filepath.Join(filepath.Dir(filename), "..", "..", "..", "internal", "storage", "queries") + queryDirectory := filepath.Join( + filepath.Dir(filename), + "..", + "..", + "..", + "internal", + "storage", + "queries", + ) issues, err := checkDirectory(queryDirectory) if err != nil { t.Fatal(err) @@ -527,3 +581,32 @@ func TestRunRejectsEmptyQueryDirectory(t *testing.T) { t.Fatalf("stderr = %q, want missing-query diagnostic", stderr.String()) } } + +func TestExecutionDMLOwnership(t *testing.T) { + for _, sql := range []string{ + `INSERT INTO agents(id) VALUES ($1)`, + `UPDATE public.agent_inputs SET state='canceled' WHERE id=$1`, + `DELETE FROM "public"."agent_events" WHERE id=$1`, + `WITH changed AS (UPDATE tool_calls SET state='ready' WHERE id=$1 RETURNING id) SELECT id FROM changed`, + `INSERT INTO agent_execution_state(agent_id) VALUES ($1) ON CONFLICT(agent_id) DO UPDATE SET stop_sequence=0`, + `INSERT /* comment */ INTO model_call_contexts(id) VALUES ($1)`, + } { + t.Run(sql, func(t *testing.T) { + source := []byte("-- name: Mutation :exec\n" + sql + ";\n") + public := checkSource("internal/storage/queries/escape.sql", source) + var found bool + for _, issue := range public { + found = found || issue.Rule == "execution-dml-ownership" + } + if !found { + t.Fatalf("missing ownership rejection: %+v", public) + } + private := checkSource("internal/storage/internal/agentexecution/queries/command.sql", source) + for _, issue := range private { + if issue.Rule == "execution-dml-ownership" { + t.Fatalf("private command rejected: %+v", issue) + } + } + }) + } +} diff --git a/tools/omnaralint/execution.go b/tools/omnaralint/execution.go new file mode 100644 index 000000000..9209f79cb --- /dev/null +++ b/tools/omnaralint/execution.go @@ -0,0 +1,145 @@ +package omnaralint + +import ( + "encoding/json" + "go/ast" + "go/constant" + "go/types" + "strings" + + pgquery "github.com/wasilibs/go-pgquery" + "golang.org/x/tools/go/analysis" +) + +const executionPackage = storagePackage + "/internal/agentexecution" + +func executionRelation(name string) bool { + switch name { + case "agents", + "agent_inputs", + "agent_events", + "agent_turns", + "model_call_contexts", + "model_outputs", + "tool_calls", + "tool_call_results", + "context_checkpoints", + "content_blocks", + "agent_interactions", + "agent_execution_state", + "agent_wakeups", + "agent_runtime_locks": + return true + } + return false +} + +func executionSQL(sql string) (bool, error) { + parsed, err := pgquery.ParseToJSON(sql) + if err != nil { + return false, err + } + var tree any + if err = json.Unmarshal([]byte(parsed), &tree); err != nil { + return false, err + } + var visit func(any) bool + visit = func(value any) bool { + switch node := value.(type) { + case []any: + for _, child := range node { + if visit(child) { + return true + } + } + case map[string]any: + for kind, child := range node { + switch kind { + case "InsertStmt", "UpdateStmt", "DeleteStmt", "MergeStmt", "CopyStmt": + statement, _ := child.(map[string]any) + relation, _ := statement["relation"].(map[string]any) + name, _ := relation["relname"].(string) + if executionRelation(name) { + return true + } + case "DoStmt", "CallStmt", "TruncateStmt", "CreateFunctionStmt", "AlterTableStmt", "DropStmt": + return true + } + if visit(child) { + return true + } + } + } + return false + } + return visit(tree), nil +} + +func checkExecutionAccess(pass *analysis.Pass, file *ast.File, isTest bool) { + if isTest || packageWithin(pass.Pkg.Path(), executionPackage) || + packageWithin(pass.Pkg.Path(), dbsqlcImport) || + packageWithin(pass.Pkg.Path(), testutilPackage) || + packageWithin(pass.Pkg.Path(), repoPackage+"/migrations") { + return + } + facade := packageWithin(pass.Pkg.Path(), storagePackage) + queryForwarder := pass.Pkg.Path() == storagePackage+"/internal/storeutil" && + strings.HasSuffix(pass.Fset.Position(file.Pos()).Filename, "/pool.go") + for _, imp := range file.Imports { + if packageWithin(strings.Trim(imp.Path.Value, `"`), "github.com/omnara-ai/omnara/internal/testutil") { + pass.Reportf(imp.Pos(), "production code cannot import execution test fixtures") + } + if strings.Trim(imp.Path.Value, `"`) == executionPackage+"/internal/executiondb" { + pass.Reportf(imp.Pos(), "executiondb imports belong to agentexecution") + } + } + ast.Inspect(file, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + if selector, ok := call.Fun.(*ast.SelectorExpr); ok && + (selector.Sel.Name == "CopyFrom" || selector.Sel.Name == "SendBatch") && + facade { + pass.Reportf(call.Pos(), "execution-capable bulk SQL belongs to agentexecution") + return true + } + signature, ok := pass.TypesInfo.TypeOf(call.Fun).(*types.Signature) + if ok && facade && + (strings.Contains(types.TypeString(signature, + nil), + "github.com/jackc/pgx/v5.Batch") || + strings.Contains(types.TypeString(signature, + nil), + "github.com/jackc/pgx/v5.CopyFromSource")) { + pass.Reportf(call.Pos(), "execution-capable bulk SQL belongs to agentexecution") + return true + } + if !ok || signature.Params().Len() < 2 || signature.Results().Len() == 0 || len(call.Args) < 2 { + return true + } + if types.TypeString(signature.Params().At(0).Type(), nil) != "context.Context" || + types.TypeString(signature.Params().At(1).Type(), nil) != "string" { + return true + } + result := types.TypeString(signature.Results().At(0).Type(), nil) + if !strings.Contains(result, "github.com/jackc/pgx/") && !strings.Contains(result, "database/sql.") { + return true + } + value := pass.TypesInfo.Types[call.Args[1]].Value + if value == nil || value.Kind() != constant.String { + if facade && !queryForwarder { + pass.Reportf(call.Pos(), "dynamic SQL cannot bypass the execution DML boundary") + } + return true + } + writes, err := executionSQL(constant.StringVal(value)) + if writes { + pass.Reportf(call.Pos(), "execution writes belong to agentexecution semantic commands") + } + if err != nil && facade { + pass.Reportf(call.Pos(), "SQL must be statically parseable outside agentexecution") + } + return true + }) +} diff --git a/tools/omnaralint/go.mod b/tools/omnaralint/go.mod index b0f66e9a2..6016741c1 100644 --- a/tools/omnaralint/go.mod +++ b/tools/omnaralint/go.mod @@ -4,10 +4,16 @@ go 1.27.2 require ( github.com/golangci/plugin-module-register v0.1.2 + github.com/wasilibs/go-pgquery v0.0.0-20260623022807-b68b28131ed3 golang.org/x/tools v0.50.0 ) require ( + github.com/pganalyze/pg_query_go/v6 v6.2.2 // indirect + github.com/tetratelabs/wazero v1.12.0 // indirect + github.com/wasilibs/wazero-helpers v0.0.0-20250123031827-cd30c44769bb // indirect golang.org/x/mod v0.41.0 // indirect golang.org/x/sync v0.23.0 // indirect + golang.org/x/sys v0.48.0 // indirect + google.golang.org/protobuf v1.36.11 // indirect ) diff --git a/tools/omnaralint/go.sum b/tools/omnaralint/go.sum index 8f5ecf09e..99d55721f 100644 --- a/tools/omnaralint/go.sum +++ b/tools/omnaralint/go.sum @@ -1,10 +1,30 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/golangci/plugin-module-register v0.1.2 h1:e5WM6PO6NIAEcij3B053CohVp3HIYbzSuP53UAYgOpg= github.com/golangci/plugin-module-register v0.1.2/go.mod h1:1+QGTsKBvAIvPvoY/os+G5eoqxWn70HYDm2uvUyGuVw= -github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= -github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/pganalyze/pg_query_go/v6 v6.2.2 h1:O0L6zMC226R82RF3X5n0Ki6HjytDsoAzuzp4ATVAHNo= +github.com/pganalyze/pg_query_go/v6 v6.2.2/go.mod h1:Cn6+j4870kJz3iYNsb0VsNG04vpSWgEvBwc590J4qD0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/tetratelabs/wazero v1.12.0 h1:DuWcpNu/FzgEXgGBDp8J1Spc+CWOvvtvVyjKlaZopYU= +github.com/tetratelabs/wazero v1.12.0/go.mod h1:LvKtzl2RqO4gyF27BiXU+nKAjcV8f38U+kP/q2vgxh0= +github.com/wasilibs/go-pgquery v0.0.0-20260623022807-b68b28131ed3 h1:XRtDsFHBMM6PNAh4EW7ChRjuZ2C9Vo19/1Gdp6G1NR0= +github.com/wasilibs/go-pgquery v0.0.0-20260623022807-b68b28131ed3/go.mod h1:ZSyYLCRbk2xPqu7lgfrDSSHm+g/7Rxk6JK4KE2cxJ3s= +github.com/wasilibs/wazero-helpers v0.0.0-20250123031827-cd30c44769bb h1:gQ+ZV4wJke/EBKYciZ2MshEouEHFuinB85dY3f5s1q8= +github.com/wasilibs/wazero-helpers v0.0.0-20250123031827-cd30c44769bb/go.mod h1:jMeV4Vpbi8osrE/pKUxRZkVaA0EX7NZN0A9/oRzgpgY= golang.org/x/mod v0.41.0 h1:qJmnOUb4YB+FsEuM3HcWucdZASCPGhsX6uljO6pog0c= golang.org/x/mod v0.41.0/go.mod h1:Ek9pY8RKWXwsWvd3rQiHYtMqkjSUV+s1Rj7j4H5Ur6o= golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk= golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= +golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo= +golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= golang.org/x/tools v0.50.0 h1:c2ifzfcuY7L90lZ2aKd8S4K2NpASF08SZx9ZuJkHmSU= golang.org/x/tools v0.50.0/go.mod h1:7ulVMw3831Mwi5EZD6RomGyffr4VFjuNYXf2BbCEAV0= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/tools/omnaralint/omnaralint.go b/tools/omnaralint/omnaralint.go index 65903c405..0468e09f8 100644 --- a/tools/omnaralint/omnaralint.go +++ b/tools/omnaralint/omnaralint.go @@ -77,15 +77,20 @@ func checkGoFile(pass *analysis.Pass, file *ast.File) { filename := slashPath(pass.Fset.Position(file.Pos()).Filename) packagePath := pass.Pkg.Path() inStorage := packageWithin(packagePath, storagePackage) - inGeneratedDBSQLC := packageWithin(packagePath, dbsqlcImport) || packageWithin(packagePath, stateDBSQLCImport) + inGeneratedDBSQLC := packageWithin(packagePath, dbsqlcImport) || + packageWithin(packagePath, stateDBSQLCImport) inInternal := packageWithin(packagePath, internalPackage) inTestutil := packageWithin(packagePath, testutilPackage) isTest := strings.HasSuffix(filename, "_test.go") + checkExecutionAccess(pass, file, isTest) for _, imp := range file.Imports { importPath := strings.Trim(imp.Path.Value, `"`) if importPath == dbsqlcImport && !inStorage { - pass.Reportf(imp.Pos(), "internal/storage/internal/dbsqlc imports must stay inside internal/storage") + pass.Reportf( + imp.Pos(), + "internal/storage/internal/dbsqlc imports must stay inside internal/storage", + ) } if forbiddenStateDBImport(packagePath, importPath) { pass.Reportf( @@ -106,7 +111,8 @@ func checkGoFile(pass *analysis.Pass, file *ast.File) { } case *ast.InterfaceType: if inGeneratedDBSQLC && !isAllowedDBSQLCInterface(pass, n) { - pass.Reportf(n.Pos(), "unexpected interface{} in generated dbsqlc; add an explicit SQL cast at the query boundary") + pass.Reportf(n.Pos(), + "unexpected interface{} in generated dbsqlc; add an explicit SQL cast at the query boundary") } case *ast.SelectorExpr: if inGeneratedDBSQLC { diff --git a/tools/omnaralint/omnaralint_test.go b/tools/omnaralint/omnaralint_test.go index cbdfd0e75..3f73e5701 100644 --- a/tools/omnaralint/omnaralint_test.go +++ b/tools/omnaralint/omnaralint_test.go @@ -1,10 +1,15 @@ package omnaralint import ( + "go/parser" + "go/token" + "go/types" "os" "path/filepath" "testing" + "golang.org/x/tools/go/analysis" + "golang.org/x/tools/go/analysis/analysistest" ) @@ -84,3 +89,65 @@ func TestStateSQLWildcardScan(t *testing.T) { t.Fatalf("wildcard issues = %v, want 2", issues) } } + +func TestExecutionSQLBoundary(t *testing.T) { + for _, sql := range []string{ + `INSERT INTO agents(id) VALUES ($1)`, + `INSERT /* nested /* comment */ */ INTO "public"."agent_inputs"(id) VALUES ($1)`, + `UPDATE tool_calls SET state='ready' WHERE id=$1`, + `WITH changed AS (DELETE FROM agent_events WHERE id=$1 RETURNING id) SELECT id FROM changed`, + `COPY agents FROM STDIN`, + `MERGE INTO agent_inputs a USING old_inputs b ON a.id=b.id WHEN MATCHED THEN DELETE`, + `TRUNCATE agents`, + `DO $$ BEGIN DELETE FROM agents; END $$`, + } { + t.Run(sql, func(t *testing.T) { + writes, err := executionSQL(sql) + if err != nil || !writes { + t.Fatalf("writes=%v error=%v", writes, err) + } + }) + } + for _, sql := range []string{`SELECT id FROM agents WHERE id=$1 FOR UPDATE`, + `SELECT 'INSERT INTO agents'`, + `UPDATE machines SET display_name=$2 WHERE id=$1`} { + writes, err := executionSQL(sql) + if err != nil || writes { + t.Fatalf("%s: writes=%v error=%v", sql, writes, err) + } + } +} + +func TestExecutionGoBoundary(t *testing.T) { + analysistest.Run( + t, + analysistest.TestData(), + Analyzer, + "github.com/omnara-ai/omnara/internal/storage/executionstore", + ) +} + +func TestExecutionPrivateImport(t *testing.T) { + for _, test := range []struct{ path, message string }{ + {executionPackage + "/internal/executiondb", "executiondb imports belong to agentexecution"}, + {testutilPackage + "/storagetest", "production code cannot import execution test fixtures"}, + } { + t.Run(test.path, func(t *testing.T) { + set := token.NewFileSet() + file, err := parser.ParseFile(set, "escape.go", "package executionstore\n import \""+test.path+"\"", 0) + if err != nil { + t.Fatal(err) + } + var messages []string + pass := &analysis.Pass{ + Fset: set, + Pkg: types.NewPackage(storagePackage+"/executionstore", "executionstore"), + Report: func(d analysis.Diagnostic) { messages = append(messages, d.Message) }, + } + checkExecutionAccess(pass, file, false) + if len(messages) != 1 || messages[0] != test.message { + t.Fatalf("diagnostics=%v", messages) + } + }) + } +} diff --git a/tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgconn/pgconn.go b/tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgconn/pgconn.go new file mode 100644 index 000000000..d264a6ac4 --- /dev/null +++ b/tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgconn/pgconn.go @@ -0,0 +1,3 @@ +package pgconn + +type CommandTag struct{} diff --git a/tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgx.go b/tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgx.go new file mode 100644 index 000000000..24d408c16 --- /dev/null +++ b/tools/omnaralint/testdata/src/github.com/jackc/pgx/v5/pgx.go @@ -0,0 +1,8 @@ +package pgx + +type Row interface{ Scan(...any) error } +type Batch struct{} +type BatchResults interface{ Close() error } + +type Identifier []string +type CopyFromSource interface{ Next() bool } diff --git a/tools/omnaralint/testdata/src/github.com/omnara-ai/omnara/internal/storage/executionstore/escapes.go b/tools/omnaralint/testdata/src/github.com/omnara-ai/omnara/internal/storage/executionstore/escapes.go new file mode 100644 index 000000000..49919fc88 --- /dev/null +++ b/tools/omnaralint/testdata/src/github.com/omnara-ai/omnara/internal/storage/executionstore/escapes.go @@ -0,0 +1,62 @@ +package executionstore + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type executor interface { + Exec(context.Context, string, ...any) (pgconn.CommandTag, error) + QueryRow(context.Context, string, ...any) pgx.Row + SendBatch(context.Context, *pgx.Batch) pgx.BatchResults + CopyFrom(context.Context, pgx.Identifier, []string, pgx.CopyFromSource) (int64, error) +} + +func escapes(ctx context.Context, db executor, sql string) { + db.Exec( // want "execution writes belong to agentexecution" + ctx, + `INSERT INTO agents(id) VALUES ($1)`, + 1, + ) + db.QueryRow( // want "execution writes belong to agentexecution" + ctx, + `WITH changed AS (DELETE FROM agent_inputs WHERE id=$1 RETURNING id) SELECT id FROM changed`, + 1, + ) + db.Exec( // want "dynamic SQL cannot bypass" + ctx, + sql, + ) + exec := db.Exec + exec( // want "execution writes belong to agentexecution" + ctx, + `UPDATE /* comment */ "public"."tool_calls" SET state='ready' WHERE id=$1`, + 1, + ) + db.SendBatch( // want "execution-capable bulk SQL belongs to agentexecution" + ctx, + &pgx.Batch{}, + ) + batch := db.SendBatch + batch( // want "execution-capable bulk SQL belongs to agentexecution" + ctx, + &pgx.Batch{}, + ) + db.CopyFrom( // want "execution-capable bulk SQL belongs to agentexecution" + ctx, + pgx.Identifier{"agents"}, + []string{"id"}, + nil, + ) + copyRows := db.CopyFrom + copyRows( // want "execution-capable bulk SQL belongs to agentexecution" + ctx, + pgx.Identifier{"agents"}, + []string{"id"}, + nil, + ) + db.Exec(ctx, `UPDATE machines SET display_name=$1 WHERE id=$2`, "machine", 1) + db.QueryRow(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, 1) +} From 11c8d6f239a3782352bd92b74884a08534520be5 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Fri, 9 Oct 2026 20:14:52 -0700 Subject: [PATCH 02/19] chore: format tests and run agent execution integration tests in CI Co-Authored-By: Claude Opus 5.5 (1M context) --- Makefile | 1 + .../tools/model_context_test_helpers_test.go | 8 ++++---- .../inbox_input_integration_test.go | 18 +++++++++--------- ...event_authority_helpers_integration_test.go | 16 ++++++++-------- 4 files changed, 22 insertions(+), 21 deletions(-) diff --git a/Makefile b/Makefile index 7463a495b..a450cd190 100644 --- a/Makefile +++ b/Makefile @@ -33,6 +33,7 @@ INTEGRATION_STORAGE_PACKAGES := \ ./internal/storage/executionstore \ ./internal/storage/identitystore \ ./internal/storage/integrationstore \ + ./internal/storage/internal/agentexecution \ ./internal/testutil/storagetest INTEGRATION_HTTPAPI_PACKAGES := \ ./internal/httpapi \ diff --git a/internal/harness/tools/model_context_test_helpers_test.go b/internal/harness/tools/model_context_test_helpers_test.go index 8ea29f504..4183562d7 100644 --- a/internal/harness/tools/model_context_test_helpers_test.go +++ b/internal/harness/tools/model_context_test_helpers_test.go @@ -41,10 +41,10 @@ func claimNormalModelCallForToolsTest( } prepared1, err := store.Execution(). PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ - ProjectID: projectID, - AgentID: agentID, - RuntimeLockID: runtime.ID, - OpeningInputIDs: openingInputIDs, + ProjectID: projectID, + AgentID: agentID, + RuntimeLockID: runtime.ID, + OpeningInputIDs: openingInputIDs, SourceModelCallContextID: sourceModelCallContextID, SourceModelOutputID: sourceModelOutputID, }) diff --git a/internal/storage/executionstore/inbox_input_integration_test.go b/internal/storage/executionstore/inbox_input_integration_test.go index f2c44620a..534be6c2a 100644 --- a/internal/storage/executionstore/inbox_input_integration_test.go +++ b/internal/storage/executionstore/inbox_input_integration_test.go @@ -53,18 +53,18 @@ func inboxInputIntegration( f.ctx, integrationstore.ConfigureIntegrationInput{ - OrgID: testOrgID, - ProjectID: testProjectID, - IntegrationID: integration.ID, - InstalledByUserID: f.user.ID, - Provider: provider, - ProviderTenantID: "123", - ProviderAccountRef: "456", - CredentialAppID: 123, + OrgID: testOrgID, + ProjectID: testProjectID, + IntegrationID: integration.ID, + InstalledByUserID: f.user.ID, + Provider: provider, + ProviderTenantID: "123", + ProviderAccountRef: "456", + CredentialAppID: 123, CredentialSecretID: secret.ID, CredentialVersionID: secret.CurrentVersionID, ExpectedSetupRevision: integration.SetupRevision, - ProviderIdentity: json.RawMessage(`{"bot_user_id":789,"bot_login":"kernel-test[bot]"}`), + ProviderIdentity: json.RawMessage(`{"bot_user_id":789,"bot_login":"kernel-test[bot]"}`), }, ) require.NoError(t, err) diff --git a/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go b/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go index d71dc7d5c..40411d991 100644 --- a/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go @@ -166,10 +166,10 @@ LIMIT 1 } prepared1, err := fixture.Store.Execution(). PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: []uuid.UUID{openingInputID}, + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: []uuid.UUID{openingInputID}, SourceModelCallContextID: work.ModelCallContextID, SourceModelOutputID: work.SourceModelOutputID, }) @@ -402,10 +402,10 @@ func claimNormalContextAtFrontierTest( } prepared2, err := fixture.Store.Execution(). PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ - ProjectID: testProjectID, - AgentID: fixture.AgentID, - RuntimeLockID: fixture.Lock.ID, - OpeningInputIDs: openingInputIDs, + ProjectID: testProjectID, + AgentID: fixture.AgentID, + RuntimeLockID: fixture.Lock.ID, + OpeningInputIDs: openingInputIDs, SourceModelCallContextID: work.ModelCallContextID, SourceModelOutputID: work.SourceModelOutputID, }) From f21e67cc9337f4beb47f4cb49c759cd9a00348a3 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Fri, 9 Oct 2026 21:12:51 -0700 Subject: [PATCH 03/19] perf: reuse execution facts within a unit and trim hot-path statements Keep each handle's execution view in memory across selections within a unit, updating it from successful receive, admission, preparation and output writes and invalidating it otherwise, so publication reloads only when a selection is stale. Move webhook eligibility capture into publication, simplify the facts query, validate single-input opening lineage without set operations, plan parent and root routing in one read, and reuse the claim's root identity. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../agent_interactions_store.go | 6 +- .../agent_runtime_lock_reaping_store.go | 6 +- .../executionstore/execution_mutation.go | 7 +- .../storage/executionstore/subagents_store.go | 50 ------- .../agentexecution/commands_failure.go | 4 +- .../commands_family_integration_test.go | 8 +- .../agentexecution/commands_inputs.go | 54 +++++++- .../commands_integration_test.go | 28 ++++ .../agentexecution/commands_models.go | 74 +++++++---- .../agentexecution/commands_outputs.go | 36 ++++- .../agentexecution/commands_runtime.go | 47 ++++--- .../internal/agentexecution/commands_tools.go | 9 +- .../storage/internal/agentexecution/handle.go | 13 +- .../internal/executiondb/batch.go | 123 ++++++++--------- .../internal/executiondb/load.sql.go | 11 ++ .../internal/executiondb/state.sql.go | 21 +++ .../storage/internal/agentexecution/load.go | 2 +- .../load_count_integration_test.go | 66 ++++++++++ .../agentexecution/load_integration_test.go | 10 ++ .../storage/internal/agentexecution/locks.go | 65 ++++++--- .../internal/agentexecution/mutation.go | 124 +++++++++++------- .../publication_integration_test.go | 8 +- .../internal/agentexecution/queries/load.sql | 38 +++--- .../agentexecution/queries/publication.sql | 12 +- .../internal/agentexecution/queries/state.sql | 3 + .../agentexecution/schema_integration_test.go | 12 ++ .../internal/agentexecution/selection_load.go | 27 +++- .../transition_count_integration_test.go | 78 +++++++++++ .../internal/dbsqlc/agent_inputs.sql.go | 13 +- .../internal/dbsqlc/agent_runtime.sql.go | 5 +- .../storage/internal/dbsqlc/subagents.sql.go | 19 --- internal/storage/queries/agent_inputs.sql | 6 +- internal/storage/queries/agent_runtime.sql | 2 +- internal/storage/queries/subagents.sql | 6 - migrations/000052_agent_execution_state.sql | 19 ++- 35 files changed, 708 insertions(+), 304 deletions(-) create mode 100644 internal/storage/internal/agentexecution/transition_count_integration_test.go diff --git a/internal/storage/executionstore/agent_interactions_store.go b/internal/storage/executionstore/agent_interactions_store.go index f3be585cc..255a4a68d 100644 --- a/internal/storage/executionstore/agent_interactions_store.go +++ b/internal/storage/executionstore/agent_interactions_store.go @@ -109,7 +109,11 @@ func (t *toolCallTransaction) createQuestionInteraction( ctx context.Context, input CreateQuestionInteractionInput, ) (AgentInteractionRecord, error) { - if err := lockAgentWithParentTx(ctx, t.unit, t.q, t.input.ProjectID, t.input.AgentID); err != nil { + plan, err := t.unit.PlanAgentFamily(ctx, t.input.ProjectID, t.input.AgentID, agentexecution.LifecycleAuthority{}) + if err == nil { + err = t.unit.LockAgents(ctx, plan) + } + if err != nil { return AgentInteractionRecord{}, err } destination, err := captureInteractionDestinationTx(ctx, t.tx, t.input.ProjectID, t.input.AgentID) diff --git a/internal/storage/executionstore/agent_runtime_lock_reaping_store.go b/internal/storage/executionstore/agent_runtime_lock_reaping_store.go index cf12a279d..00e646cb8 100644 --- a/internal/storage/executionstore/agent_runtime_lock_reaping_store.go +++ b/internal/storage/executionstore/agent_runtime_lock_reaping_store.go @@ -97,7 +97,11 @@ func reapExpiredAgentRuntimeLockTx( unit *agentexecution.Unit, projectID, agentID, runtimeLockID uuid.UUID, ) (bool, error) { - locked, err := tryLockAgentWithParentTx(ctx, unit, dbsqlc.New(unit.DB()), projectID, agentID) + plan, err := unit.PlanAgentFamily(ctx, projectID, agentID, agentexecution.LifecycleAuthority{}) + var locked bool + if err == nil { + locked, err = unit.TryLockAgents(ctx, plan) + } if errors.Is(err, storeerr.ErrNotFound) { return false, nil } diff --git a/internal/storage/executionstore/execution_mutation.go b/internal/storage/executionstore/execution_mutation.go index 90d839600..7cb2d58f7 100644 --- a/internal/storage/executionstore/execution_mutation.go +++ b/internal/storage/executionstore/execution_mutation.go @@ -5,7 +5,6 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) func beginExecution( @@ -17,7 +16,11 @@ func beginExecution( if err != nil { return nil, nil, err } - if err = lockAgentWithParentTx(ctx, unit, dbsqlc.New(unit.DB()), projectID, agentID); err != nil { + plan, err := unit.PlanAgentFamily(ctx, projectID, agentID, agentexecution.LifecycleAuthority{}) + if err == nil { + err = unit.LockAgents(ctx, plan) + } + if err != nil { _ = unit.Rollback(ctx) return nil, nil, err } diff --git a/internal/storage/executionstore/subagents_store.go b/internal/storage/executionstore/subagents_store.go index 7761c0236..f79267f5c 100644 --- a/internal/storage/executionstore/subagents_store.go +++ b/internal/storage/executionstore/subagents_store.go @@ -10,7 +10,6 @@ import ( "time" "github.com/google/uuid" - "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/publicid" "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" @@ -326,55 +325,6 @@ func resolveSubagentTx( ) } -func lockAgentWithParentTx( - ctx context.Context, unit *agentexecution.Unit, qtx *dbsqlc.Queries, projectID, agentID uuid.UUID, -) error { - refs, err := agentWithParentLockRefsTx(ctx, qtx, projectID, agentID) - if err != nil { - return err - } - return unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}) -} - -func tryLockAgentWithParentTx( - ctx context.Context, - unit *agentexecution.Unit, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) (bool, error) { - refs, err := agentWithParentLockRefsTx(ctx, qtx, projectID, agentID) - if err != nil { - return false, err - } - plan, err := unit.PlanAgents(ctx, refs, agentexecution.LifecycleAuthority{}) - if err != nil { - return false, err - } - return unit.TryLockAgents(ctx, plan) -} - -func agentWithParentLockRefsTx( - ctx context.Context, - qtx *dbsqlc.Queries, - projectID, agentID uuid.UUID, -) ([]lifecyclelock.AgentRef, error) { - parentID, err := qtx.GetAgentParentID( - ctx, - dbsqlc.GetAgentParentIDParams{ProjectID: projectID, ID: agentID}, - ) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return nil, storeerr.ErrNotFound - } - return nil, fmt.Errorf("load agent parent: %w", err) - } - refs := []lifecyclelock.AgentRef{{ProjectID: projectID, AgentID: agentID}} - if parentID != nil { - refs = append(refs, lifecyclelock.AgentRef{ProjectID: projectID, AgentID: *parentID}) - } - return refs, nil -} - func textInputContentBlocks(text string) ([]CreateContentBlockInput, json.RawMessage, error) { contentBlocks := []CreateContentBlockInput{{ Ordinal: 0, diff --git a/internal/storage/internal/agentexecution/commands_failure.go b/internal/storage/internal/agentexecution/commands_failure.go index ff450e34f..41b8a2deb 100644 --- a/internal/storage/internal/agentexecution/commands_failure.go +++ b/internal/storage/internal/agentexecution/commands_failure.go @@ -245,9 +245,11 @@ func (h *Handle) publishTerminalFailure(ctx context.Context, m *executionMutatio if err = h.advanceTurn(ctx, event, true); err != nil { return false, err } + sourceContext := attemptRecord(row).Context + sourceContext.State = ContextFailed if err = applyAcceptedOutput(m, snapshot, - attemptRecord(row).Context, + sourceContext, output.ID, event, false, diff --git a/internal/storage/internal/agentexecution/commands_family_integration_test.go b/internal/storage/internal/agentexecution/commands_family_integration_test.go index 64b7b0698..523b7d3b2 100644 --- a/internal/storage/internal/agentexecution/commands_family_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_family_integration_test.go @@ -57,8 +57,8 @@ func TestExecutionCreatedAgentPublicationAndSavepoint(t *testing.T) { require.ErrorIs(t, err, agentexecution.ErrUnitClosed) statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) require.NoError(t, err) - require.Equal(t, 2, statements) - require.Equal(t, 2, roundTrips) + require.Equal(t, 1, statements) + require.Equal(t, 1, roundTrips) require.NoError( t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM agent_execution_state`).Scan(&heads), @@ -175,8 +175,8 @@ func TestExecutionParentEffects(t *testing.T) { require.Equal(t, 1, inputs) statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) require.NoError(t, err) - require.Equal(t, 4, statements) - require.Equal(t, 2, roundTrips) + require.Equal(t, 2, statements) + require.Equal(t, 1, roundTrips) }, ) } diff --git a/internal/storage/internal/agentexecution/commands_inputs.go b/internal/storage/internal/agentexecution/commands_inputs.go index 2b789cb34..076206476 100644 --- a/internal/storage/internal/agentexecution/commands_inputs.go +++ b/internal/storage/internal/agentexecution/commands_inputs.go @@ -99,7 +99,19 @@ func (h *Handle) ReceiveContent(ctx context.Context, input ReceiveContentInput) if err = h.writeContent(ctx, "agent_input", row.ID, input.Content); err != nil { return ReceivedContent{}, err } - m.changed() + if m.loaded != nil { + m.loaded.databaseNow = row.QueuedAt + if input.DeliveryMode == "steering" { + m.loaded.View.Inputs.Steering = true + } else { + m.loaded.View.Inputs.Queued = true + } + if err := m.updated(*m.loaded); err != nil { + return ReceivedContent{}, err + } + } else { + m.changed() + } return ReceivedContent{ID: row.ID, Created: true}, nil }) } @@ -170,7 +182,47 @@ func (h *Handle) admitInputs( AgentID: h.route.AgentID, ID: result.TurnID, EventID: last.ID}); err != nil { return InputAdmission{}, err } + loaded := m.loaded m.openTurn(result.TurnID) + if loaded == nil { + return result, nil + } + opening, err := q.CaptureExecutionOpening(ctx, h.unit.DB(), executiondb.CaptureExecutionOpeningParams{ + AgentID: h.route.AgentID, ProjectID: h.route.ProjectID, TurnID: result.TurnID, + Watermark: last.Sequence, StopSequence: m.head.StopSequence, + }) + if err != nil { + return InputAdmission{}, err + } + if len(opening) == 0 { + return InputAdmission{}, ErrInvalidState + } + first := result.Inputs[0].Event + turn := &Turn{ID: result.TurnID, FirstOpeningSequence: first.Sequence, LastOpeningSequence: last.Sequence, + FirstContentSequence: first.Sequence, LatestSemantic: last.EventBoundary, + InitialOpening: Opening{EventSequence: opening[0].Sequence}, InitialReadyAt: opening[0].CreatedAt} + for _, row := range opening { + turn.InitialOpening.InputIDs = append(turn.InitialOpening.InputIDs, row.ID) + } + loaded.View = ExecutionView{AgentID: h.route.AgentID, State: loaded.View.State, Turn: turn} + admitted := make(map[uuid.UUID]bool, len(result.Inputs)) + for _, input := range result.Inputs { + admitted[input.ID] = true + } + for _, row := range rows { + if admitted[row.ID] { + continue + } + if row.DeliveryMode == "steering" { + loaded.View.Inputs.Steering = true + } else { + loaded.View.Inputs.Queued = true + } + } + loaded.databaseNow = last.Time + if err := m.updated(*loaded); err != nil { + return InputAdmission{}, err + } return result, nil } diff --git a/internal/storage/internal/agentexecution/commands_integration_test.go b/internal/storage/internal/agentexecution/commands_integration_test.go index c066538ac..bb9242a3b 100644 --- a/internal/storage/internal/agentexecution/commands_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_integration_test.go @@ -42,6 +42,9 @@ func assertCommandState( t.Helper() got, err := h.LoadExecution(t.Context()) require.NoError(t, err) + reloaded, err := agentexecution.ReloadExecution(t.Context(), h) + require.NoError(t, err) + require.Equal(t, reloaded.View, got.View) rebuilt, err := h.ReconstructExecution(t.Context()) require.NoError(t, err) left, right := got.Head, rebuilt.Head @@ -221,6 +224,9 @@ func TestExecutionCommandSavepointAndReplay(t *testing.T) { assertCommandState(t, u, h) return sentinel }), sentinel) + _, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + require.NoError(t, err) + require.Zero(t, count) got := assertCommandState(t, u, h) require.Equal(t, expected.Head.CurrentTurnID, got.Head.CurrentTurnID) require.Equal(t, agentexecution.AdmitOneQueued, got.Selection.Admission) @@ -957,3 +963,25 @@ func TestExecutionContentRejectsDatabaseUnsafeValues(t *testing.T) { }) } } + +func TestExecutionPreparationOwnsCachedLineage(t *testing.T) { + f, a, lease := commandFixture(t) + u, h := f.handle(t, a) + input := receiveCommand(t, h, "queued", "initial") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + selected, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + prepared, err := h.PrepareModel(t.Context(), agentexecution.PrepareModelInput{ + RuntimeLockID: lease, Selected: *selected.Selection.Model, + }) + require.NoError(t, err) + selected.Selection.Model.Opening.InputIDs[0] = uuid.New() + prepared.Context.Opening.InputIDs[0] = uuid.New() + snapshot, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + require.NoError(t, err) + require.Zero(t, count) + require.Equal(t, []uuid.UUID{input.ID}, snapshot.View.NormalContext.Opening.InputIDs) + snapshot.View.NormalContext.Opening.InputIDs[0] = uuid.New() + assertCommandState(t, u, h) +} diff --git a/internal/storage/internal/agentexecution/commands_models.go b/internal/storage/internal/agentexecution/commands_models.go index 1f9834c09..d3f07bf74 100644 --- a/internal/storage/internal/agentexecution/commands_models.go +++ b/internal/storage/internal/agentexecution/commands_models.go @@ -4,6 +4,7 @@ import ( "context" "errors" "slices" + "time" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -109,34 +110,36 @@ func (h *Handle) PrepareModel(ctx context.Context, input PrepareModelInput) (Pre desired.Context.RetryAt = nil desired.RuntimeLockID = input.RuntimeLockID } - id, err := q.LatestExecutionAttempt( - ctx, - h.unit.DB(), - executiondb.LatestExecutionAttemptParams{AgentID: h.route.AgentID, - Operation: string(desired.Context.Operation), - Watermark: desired.Context.InputEventSequence, - SourceEnd: optionalSequence(desired.Context.SourceEventSequenceEnd)}, - ) - if err == nil { - existing, err := q.ReadExecutionAttempt( + if desired.Context.InputEventSequence <= m.head.MaxContextInputSequence { + id, err := q.LatestExecutionAttempt( ctx, h.unit.DB(), - executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: id}, + executiondb.LatestExecutionAttemptParams{AgentID: h.route.AgentID, + Operation: string(desired.Context.Operation), + Watermark: desired.Context.InputEventSequence, + SourceEnd: optionalSequence(desired.Context.SourceEventSequenceEnd)}, ) - if err != nil { - return PreparedModel{}, err - } - record := attemptRecord(existing) - if record.Context.Attempt >= desired.Context.Attempt { - if record.Context.TurnID != input.Selected.TurnID || record.ConfigID != desired.ConfigID || - !slices.Equal(record.Context.Opening.InputIDs, input.Selected.Opening.InputIDs) || - record.Context.Opening.EventSequence != input.Selected.Opening.EventSequence { - return PreparedModel{}, storeerr.ErrIdempotencyConflict + if err == nil { + existing, err := q.ReadExecutionAttempt( + ctx, + h.unit.DB(), + executiondb.ReadExecutionAttemptParams{AgentID: h.route.AgentID, ID: id}, + ) + if err != nil { + return PreparedModel{}, err + } + record := attemptRecord(existing) + if record.Context.Attempt >= desired.Context.Attempt { + if record.Context.TurnID != input.Selected.TurnID || record.ConfigID != desired.ConfigID || + !slices.Equal(record.Context.Opening.InputIDs, input.Selected.Opening.InputIDs) || + record.Context.Opening.EventSequence != input.Selected.Opening.EventSequence { + return PreparedModel{}, storeerr.ErrIdempotencyConflict + } + return record, nil } - return record, nil + } else if !errors.Is(err, pgx.ErrNoRows) { + return PreparedModel{}, err } - } else if !errors.Is(err, pgx.ErrNoRows) { - return PreparedModel{}, err } snapshot, err := h.LoadExecution(ctx) if err != nil { @@ -167,7 +170,7 @@ func (h *Handle) createAttempt( if err != nil { return PreparedModel{}, err } - c := desired.Context + c := cloneModel(desired.Context) _, err = q.CreateExecutionAttempt( ctx, h.unit.DB(), @@ -207,7 +210,28 @@ func (h *Handle) createAttempt( } else { m.head.CompactionContextID = c.ID } - m.changed() + if c.Operation == OperationNormal { + snapshot.View.NormalContext = &c + snapshot.View.CompactionContext = nil + if m.head.PendingConfigInputID == uuid.Nil { + snapshot.View.Config = nil + } + if m.head.PendingCheckpointID == uuid.Nil { + snapshot.View.Checkpoint = nil + } + if m.head.PendingOutputLimitID == uuid.Nil { + snapshot.View.OutputLimit = nil + } + } else { + snapshot.View.CompactionContext = &c + } + if snapshot.View.Turn != nil && m.head.MaxContextInputSequence >= snapshot.View.Turn.FirstOpeningSequence { + snapshot.View.Turn.InitialOpening = Opening{} + snapshot.View.Turn.InitialReadyAt = time.Time{} + } + if err := m.updated(snapshot); err != nil { + return PreparedModel{}, err + } desired.RevisionID = revision.CurrentRevisionID desired.Created = true desired.Claimed = revision.Allowed diff --git a/internal/storage/internal/agentexecution/commands_outputs.go b/internal/storage/internal/agentexecution/commands_outputs.go index 43e0e5105..cc6bf17eb 100644 --- a/internal/storage/internal/agentexecution/commands_outputs.go +++ b/internal/storage/internal/agentexecution/commands_outputs.go @@ -192,15 +192,23 @@ func (h *Handle) AcceptOutput(ctx context.Context, input AcceptOutputInput) (Acc ); err != nil { return AcceptedOutput{}, err } + sourceContext := attemptRecord(source).Context + sourceContext.State = ContextSucceeded if err = applyAcceptedOutput(m, snapshot, - attemptRecord(source).Context, + sourceContext, output.ID, event, len(proposals) > 0, input.Response.Normalized.StopReason == modelenvelope.StopReasonMaxTokens); err != nil { return AcceptedOutput{}, err } + for _, call := range proposals { + if call.Error != "" { + m.changed() + break + } + } if !input.Response.HasToolCalls() { kind := map[modelenvelope.StopReason]string{modelenvelope.StopReasonEndTurn: "result", modelenvelope.StopReasonRefusal: "refused", @@ -327,9 +335,20 @@ func applyAcceptedOutput( tools, limit bool, ) error { m.head.AnsweredThroughSequence = max(m.head.AnsweredThroughSequence, event.Sequence) - m.changed() + if snapshot.View.NormalContext != nil && snapshot.View.NormalContext.ID == source.ID { + snapshot.View.NormalContext = &source + } + if snapshot.View.CompactionContext != nil && snapshot.View.CompactionContext.ID == source.ID { + snapshot.View.CompactionContext = &source + } + if snapshot.View.ToolBatch != nil && snapshot.View.ToolBatch.Output.Context.ID == source.ID { + snapshot.View.ToolBatch.Output.Context = source + } + if snapshot.View.OutputLimit != nil && snapshot.View.OutputLimit.Context.ID == source.ID { + snapshot.View.OutputLimit.Context = source + } if source.TurnID != m.head.CurrentTurnID || event.Sequence < m.head.StopSequence { - return nil + return m.updated(snapshot) } if snapshot.View.ToolBatch != nil { completion := snapshot.View.ToolBatch.Completion @@ -337,14 +356,23 @@ func applyAcceptedOutput( return errors.New("accepted model output does not cover pending tool results") } m.head.PendingToolOutputID = uuid.Nil + snapshot.View.ToolBatch = nil } m.head.PendingOutputLimitID = uuid.Nil + snapshot.View.OutputLimit = nil + output := ModelOutput{ID: id, EventID: event.ID, Event: event.EventBoundary, Context: source} if tools { m.head.PendingToolOutputID = id + snapshot.View.ToolBatch = &ToolBatch{Output: output, HasIncomplete: true, HasRunnable: true} } else if limit { m.head.PendingOutputLimitID = id + snapshot.View.OutputLimit = &output + } + if !tools && snapshot.View.Turn != nil { + snapshot.View.Turn.LatestSemantic = event.EventBoundary } - return nil + snapshot.databaseNow = event.Time + return m.updated(snapshot) } func (h *Handle) rejectProposal(ctx context.Context, m *executionMutation, call outputProposal) error { diff --git a/internal/storage/internal/agentexecution/commands_runtime.go b/internal/storage/internal/agentexecution/commands_runtime.go index e2ef206f7..98494b1cd 100644 --- a/internal/storage/internal/agentexecution/commands_runtime.go +++ b/internal/storage/internal/agentexecution/commands_runtime.go @@ -47,7 +47,7 @@ func (h *Handle) Claim( return OwnedWork{}, err } if snapshot.Selection.Work == WorkNone { - m.changed() + m.dirty = true return OwnedWork{Selection: snapshot.Selection}, nil } result := OwnedWork{} @@ -78,7 +78,8 @@ func (h *Handle) Claim( return OwnedWork{}, err } h.unit.runtimeLocked = true - m.changed() + h.lockedRuntimeID = result.RuntimeID + m.dirty = true if err = q.ConsumeExecutionWakeup(ctx, h.unit.DB(), executiondb.ConsumeExecutionWakeupParams{AgentID: h.route.AgentID}); err != nil { @@ -155,7 +156,7 @@ func (h *Handle) Advance( return OwnedWork{}, err } } - m.changed() + m.dirty = true return result, nil }) } @@ -181,20 +182,25 @@ func (h *Handle) lockRuntime( id uuid.UUID, skip bool, ) (executiondb.ReadExecutionRuntimeRow, error) { + if id == uuid.Nil { + return executiondb.ReadExecutionRuntimeRow{}, storeerr.ErrRuntimeLockInactive + } q := executiondb.New() var err error - if skip { - _, err = q.TryLockExecutionRuntime( - ctx, - h.unit.DB(), - executiondb.TryLockExecutionRuntimeParams{AgentID: h.route.AgentID, ID: id}, - ) - } else { - _, - err = q.LockExecutionRuntime(ctx, - h.unit.DB(), - executiondb.LockExecutionRuntimeParams{AgentID: h.route.AgentID, - ID: id}) + if h.lockedRuntimeID != id { + if skip { + _, err = q.TryLockExecutionRuntime( + ctx, + h.unit.DB(), + executiondb.TryLockExecutionRuntimeParams{AgentID: h.route.AgentID, ID: id}, + ) + } else { + _, + err = q.LockExecutionRuntime(ctx, + h.unit.DB(), + executiondb.LockExecutionRuntimeParams{AgentID: h.route.AgentID, + ID: id}) + } } if errors.Is(err, pgx.ErrNoRows) { return executiondb.ReadExecutionRuntimeRow{}, storeerr.ErrRuntimeLockInactive @@ -203,6 +209,7 @@ func (h *Handle) lockRuntime( return executiondb.ReadExecutionRuntimeRow{}, err } h.unit.runtimeLocked = true + h.lockedRuntimeID = id return q.ReadExecutionRuntime( ctx, h.unit.DB(), @@ -417,6 +424,12 @@ func (h *Handle) releaseRuntime( if n != 1 { return storeerr.ErrRuntimeLockInactive } - m.changed() - return nil + h.lockedRuntimeID = uuid.Nil + if m.loaded == nil { + m.changed() + return nil + } + m.loaded.databaseNow = runtime.DatabaseNow + m.dirty = true + return m.selectLoaded() } diff --git a/internal/storage/internal/agentexecution/commands_tools.go b/internal/storage/internal/agentexecution/commands_tools.go index ddc66c0a2..b7c2fa4e9 100644 --- a/internal/storage/internal/agentexecution/commands_tools.go +++ b/internal/storage/internal/agentexecution/commands_tools.go @@ -54,7 +54,14 @@ func (h *Handle) transitionTool( if n != 1 { return storeerr.ErrStateTransitionConflict } - m.changed() + runnable := func(state string) bool { + return state == "awaiting_authorization" || state == "ready" && (row.Type == "built_in" || row.Type == "mcp") + } + if (row.State == "completed") != (state == "completed") || runnable(row.State) != runnable(state) { + m.changed() + } else { + m.dirty = true + } return nil } diff --git a/internal/storage/internal/agentexecution/handle.go b/internal/storage/internal/agentexecution/handle.go index e88490c2f..ce5d41c59 100644 --- a/internal/storage/internal/agentexecution/handle.go +++ b/internal/storage/internal/agentexecution/handle.go @@ -4,14 +4,17 @@ import ( "fmt" "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) type Handle struct { - unit *Unit - route AgentRoute - valid bool - created bool - mutation *executionMutation + unit *Unit + route AgentRoute + valid bool + created bool + mutation *executionMutation + identity *dbsqlc.LockAgentInProjectRow + lockedRuntimeID uuid.UUID } func (u *Unit) Agent(route AgentRoute) (*Handle, error) { diff --git a/internal/storage/internal/agentexecution/internal/executiondb/batch.go b/internal/storage/internal/agentexecution/internal/executiondb/batch.go index 4614c980c..4115d5f8a 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/batch.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/batch.go @@ -27,38 +27,40 @@ WITH scope AS MATERIALIZED ( AND i.delivery_mode='steering' AND i.state='received') AS steering, EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' AND i.delivery_mode='queued' AND i.state='received') AS queued, - coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event - ORDER BY x.sequence LIMIT 1),0)::bigint AS first_opening_sequence, - coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event - ORDER BY x.sequence DESC LIMIT 1),0)::bigint AS last_opening_sequence, - coalesce((SELECT x.sequence FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id - WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event AND i.input_kind='content' - ORDER BY x.sequence LIMIT 1),0)::bigint AS first_content_sequence - FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=$3 + coalesce(o.first_opening_sequence,0)::bigint AS first_opening_sequence, + coalesce(o.last_opening_sequence,0)::bigint AS last_opening_sequence, + coalesce(o.first_content_sequence,0)::bigint AS first_content_sequence + FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=$1 LEFT JOIN agent_events e ON e.agent_id=a.id AND e.id=t.latest_semantic_event_id - WHERE a.id=$1 AND a.project_id=$4 + LEFT JOIN LATERAL ( + SELECT min(x.sequence) AS first_opening_sequence,max(x.sequence) AS last_opening_sequence, + min(x.sequence) FILTER (WHERE i.input_kind='content') AS first_content_sequence + FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id + WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event + ) o ON true + WHERE a.id=$2 AND a.project_id=$3 ), outputs AS MATERIALIZED ( SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id - WHERE o.agent_id=$1 AND o.id IN ($5::uuid,$6::uuid) + WHERE o.agent_id=$2 AND o.id IN ($4::uuid,$5::uuid) ), batch AS MATERIALIZED ( - SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 - AND c.model_output_id=$5) AS has_tools, - EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 - AND c.model_output_id=$5 AND c.state<>'completed') AS incomplete, - EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 - AND c.model_output_id=$5 AND + SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$2 + AND c.model_output_id=$4) AS has_tools, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$2 + AND c.model_output_id=$4 AND c.state<>'completed') AS incomplete, + EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$2 + AND c.model_output_id=$4 AND (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable - WHERE $5::uuid IS NOT NULL + WHERE $4::uuid IS NOT NULL ), boundary AS MATERIALIZED ( - SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$1 + SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$2 AND e.event_kind='model_output' AND e.sequence<=(SELECT last_opening_sequence FROM scope) - ORDER BY e.sequence DESC LIMIT 1),0),$7::bigint)::bigint AS sequence - FROM scope WHERE scope.state='active' AND $8::bigintb.sequence AND e.sequence<=(SELECT last_opening_sequence FROM scope) JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id @@ -66,7 +68,7 @@ WITH scope AS MATERIALIZED ( SELECT id,sequence,created_at FROM unanswered UNION ALL SELECT i.id,e.sequence,e.created_at FROM boundary b - JOIN agent_events e ON e.agent_id=$1 AND e.turn_id=$3 AND e.is_opening_event + JOIN agent_events e ON e.agent_id=$2 AND e.turn_id=$1 AND e.is_opening_event AND e.sequence<=(SELECT last_opening_sequence FROM scope) JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id @@ -99,9 +101,9 @@ SELECT 'context'::text AS kind, 'epoch'::timestamptz AS database_now, false AS steering,false AS queued, 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM model_call_contexts c WHERE c.agent_id=$1 AND c.id IN ( - SELECT $9::uuid UNION SELECT $10::uuid - UNION SELECT model_call_context_id FROM outputs) +FROM model_call_contexts c WHERE c.agent_id=$2 AND c.id=ANY( + ARRAY[$8::uuid,$9::uuid] + || ARRAY(SELECT model_call_context_id FROM outputs)) UNION ALL SELECT 'output'::text AS kind, o.id AS id, @@ -164,7 +166,7 @@ FROM batch b LEFT JOIN LATERAL ( max(e.sequence)::bigint AS last_sequence,max(r.completed_at)::timestamptz AS ready_at FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id - WHERE c.agent_id=$1 AND c.model_output_id=$5 AND NOT b.incomplete + WHERE c.agent_id=$2 AND c.model_output_id=$4 AND NOT b.incomplete ) r ON NOT b.incomplete UNION ALL SELECT 'config'::text AS kind, @@ -195,7 +197,7 @@ SELECT 'config'::text AS kind, false AS steering,false AS queued, 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id -WHERE i.agent_id=$1 AND i.id=$11 AND i.input_kind='config_change' AND i.state='resolved' +WHERE i.agent_id=$2 AND i.id=$10 AND i.input_kind='config_change' AND i.state='resolved' UNION ALL SELECT 'checkpoint'::text AS kind, c.id AS id, @@ -225,7 +227,7 @@ SELECT 'checkpoint'::text AS kind, false AS steering,false AS queued, 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id -WHERE c.agent_id=$1 AND c.id=$12 +WHERE c.agent_id=$2 AND c.id=$11 UNION ALL SELECT 'opening'::text AS kind, o.id AS id, @@ -266,17 +268,8 @@ FROM scope s SELECT f.kind,f.id,f.turn_id,f.context_id,f.event_id,f.sequence,f.event_time,f.operation,f.attempt, f.watermark,f.source_end,f.state,f.recovery,f.retry_at,f.opening_ids,f.opening_sequence, f.has_tools,f.incomplete,f.runnable,f.calls,f.results,f.events,f.stop_reason,f.root_agent_id, - f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence, - coalesce(w.org_id,'00000000-0000-0000-0000-000000000000'::uuid)::uuid AS webhook_org_id, - coalesce(w.events,'[]'::jsonb)::jsonb AS webhook_events -FROM facts f LEFT JOIN LATERAL ( - SELECT p.org_id,c.compiled_definition->'event_webhook'->'events' AS events - FROM agents a JOIN agent_configs c ON c.project_id=a.project_id AND c.id=a.current_config_id - JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id - WHERE a.id=$1 AND p.deleted_at IS NULL AND o.deleted_at IS NULL - AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' - AND $2::boolean AND f.kind='scope' -) w ON true ORDER BY f.kind,f.sequence + f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence +FROM facts f ORDER BY f.kind,f.sequence ` type LoadExecutionFactsBatchResults struct { @@ -286,9 +279,8 @@ type LoadExecutionFactsBatchResults struct { } type LoadExecutionFactsParams struct { - AgentID uuid.UUID - CaptureWebhook bool TurnID *uuid.UUID + AgentID uuid.UUID ProjectID uuid.UUID ToolOutputID *uuid.UUID LimitOutputID *uuid.UUID @@ -331,17 +323,14 @@ type LoadExecutionFactsRow struct { FirstOpeningSequence int64 LastOpeningSequence int64 FirstContentSequence int64 - WebhookOrgID uuid.UUID - WebhookEvents json.RawMessage } func (q *Queries) LoadExecutionFacts(ctx context.Context, db DBTX, arg []LoadExecutionFactsParams) *LoadExecutionFactsBatchResults { batch := &pgx.Batch{} for _, a := range arg { vals := []interface{}{ - a.AgentID, - a.CaptureWebhook, a.TurnID, + a.AgentID, a.ProjectID, a.ToolOutputID, a.LimitOutputID, @@ -407,8 +396,6 @@ func (b *LoadExecutionFactsBatchResults) Query(f func(int, []LoadExecutionFactsR &i.FirstOpeningSequence, &i.LastOpeningSequence, &i.FirstContentSequence, - &i.WebhookOrgID, - &i.WebhookEvents, ); err != nil { return err } @@ -429,12 +416,12 @@ func (b *LoadExecutionFactsBatchResults) Close() error { const publishExecution = `-- name: PublishExecution :batchone WITH scope AS MATERIALIZED ( - SELECT a.id AS agent_id,$1::timestamptz AS logical_ready_at - FROM agents a WHERE a.id=$2 AND ($3::boolean OR + SELECT a.id AS agent_id,$2::timestamptz AS logical_ready_at + FROM agents a WHERE a.id=$3 AND ($4::boolean OR EXISTS(SELECT 1 FROM agent_execution_state h WHERE h.agent_id=a.id)) ), updated AS ( INSERT INTO agent_execution_state AS h(agent_id,current_turn_id,stop_sequence,answered_through_sequence,max_normal_input_sequence,max_context_input_sequence,normal_context_id,compaction_context_id,pending_tool_output_id,pending_output_limit_id,pending_config_input_id,pending_checkpoint_id,turn_continuable,incomplete_tools,logical_ready_at) -SELECT scope.agent_id,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,scope.logical_ready_at FROM scope +SELECT scope.agent_id,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,scope.logical_ready_at FROM scope ON CONFLICT(agent_id) DO UPDATE SET current_turn_id=excluded.current_turn_id, stop_sequence=excluded.stop_sequence, @@ -454,10 +441,10 @@ WHERE (h.current_turn_id,h.stop_sequence,h.answered_through_sequence,h.max_norma (excluded.current_turn_id,excluded.stop_sequence,excluded.answered_through_sequence,excluded.max_normal_input_sequence,excluded.max_context_input_sequence,excluded.normal_context_id,excluded.compaction_context_id,excluded.pending_tool_output_id,excluded.pending_output_limit_id,excluded.pending_config_input_id,excluded.pending_checkpoint_id,excluded.turn_continuable,excluded.incomplete_tools,excluded.logical_ready_at) RETURNING h.agent_id), desired AS MATERIALIZED ( SELECT a.id,h.logical_ready_at AS ready_at FROM scope h JOIN agents a ON a.id=h.agent_id - WHERE a.id=$2 AND a.state='active' AND h.logical_ready_at IS NOT NULL + WHERE a.id=$3 AND a.state='active' AND h.logical_ready_at IS NOT NULL AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id) ), removed AS ( - DELETE FROM agent_wakeups WHERE agent_id=$2 AND NOT EXISTS(SELECT 1 FROM desired) + DELETE FROM agent_wakeups WHERE agent_id=$3 AND NOT EXISTS(SELECT 1 FROM desired) ), published AS ( INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) SELECT id,ready_at,statement_timestamp() FROM desired @@ -469,7 +456,17 @@ WHERE agent_wakeups.ready_at IS DISTINCT FROM CASE WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END RETURNING agent_id -) SELECT agent_id FROM scope +) SELECT scope.agent_id, + coalesce(w.org_id,'00000000-0000-0000-0000-000000000000'::uuid)::uuid AS webhook_org_id, + coalesce(w.events,'[]'::jsonb)::jsonb AS webhook_events +FROM scope LEFT JOIN LATERAL ( + SELECT p.org_id,c.compiled_definition->'event_webhook'->'events' AS events + FROM agents a JOIN agent_configs c ON c.project_id=a.project_id AND c.id=a.current_config_id + JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id + WHERE a.id=scope.agent_id AND p.deleted_at IS NULL AND o.deleted_at IS NULL + AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' + AND $1::boolean +) w ON true ` type PublishExecutionBatchResults struct { @@ -479,6 +476,7 @@ type PublishExecutionBatchResults struct { } type PublishExecutionParams struct { + CaptureWebhook bool LogicalReadyAt *time.Time AgentID uuid.UUID InsertHead bool @@ -497,10 +495,17 @@ type PublishExecutionParams struct { IncompleteTools bool } +type PublishExecutionRow struct { + AgentID uuid.UUID + WebhookOrgID uuid.UUID + WebhookEvents json.RawMessage +} + func (q *Queries) PublishExecution(ctx context.Context, db DBTX, arg []PublishExecutionParams) *PublishExecutionBatchResults { batch := &pgx.Batch{} for _, a := range arg { vals := []interface{}{ + a.CaptureWebhook, a.LogicalReadyAt, a.AgentID, a.InsertHead, @@ -524,20 +529,20 @@ func (q *Queries) PublishExecution(ctx context.Context, db DBTX, arg []PublishEx return &PublishExecutionBatchResults{br, len(arg), false} } -func (b *PublishExecutionBatchResults) QueryRow(f func(int, uuid.UUID, error)) { +func (b *PublishExecutionBatchResults) QueryRow(f func(int, PublishExecutionRow, error)) { defer b.br.Close() for t := 0; t < b.tot; t++ { - var agent_id uuid.UUID + var i PublishExecutionRow if b.closed { if f != nil { - f(t, agent_id, ErrBatchAlreadyClosed) + f(t, i, ErrBatchAlreadyClosed) } continue } row := b.br.QueryRow() - err := row.Scan(&agent_id) + err := row.Scan(&i.AgentID, &i.WebhookOrgID, &i.WebhookEvents) if f != nil { - f(t, agent_id, err) + f(t, i, err) } } } diff --git a/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go index ca4e5cb82..13813c9b4 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go @@ -12,6 +12,17 @@ import ( "github.com/google/uuid" ) +const executionDatabaseTime = `-- name: ExecutionDatabaseTime :one +SELECT statement_timestamp()::timestamptz +` + +func (q *Queries) ExecutionDatabaseTime(ctx context.Context, db DBTX) (time.Time, error) { + row := db.QueryRow(ctx, executionDatabaseTime) + var column_1 time.Time + err := row.Scan(&column_1) + return column_1, err +} + const loadExecutionBase = `-- name: LoadExecutionBase :one SELECT coalesce(h.agent_id,a.id)::uuid AS agent_id,h.current_turn_id, coalesce(h.stop_sequence,0)::bigint AS stop_sequence, diff --git a/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go index d17182f11..3ce9c759e 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go @@ -40,6 +40,27 @@ func (q *Queries) LoadExecutionBoundaries(ctx context.Context, db DBTX, arg Load return i, err } +const loadExecutionFamily = `-- name: LoadExecutionFamily :one +SELECT root_agent_id,parent_agent_id FROM agents WHERE project_id=$1 AND id=$2 +` + +type LoadExecutionFamilyParams struct { + ProjectID uuid.UUID + ID uuid.UUID +} + +type LoadExecutionFamilyRow struct { + RootAgentID uuid.UUID + ParentAgentID *uuid.UUID +} + +func (q *Queries) LoadExecutionFamily(ctx context.Context, db DBTX, arg LoadExecutionFamilyParams) (LoadExecutionFamilyRow, error) { + row := db.QueryRow(ctx, loadExecutionFamily, arg.ProjectID, arg.ID) + var i LoadExecutionFamilyRow + err := row.Scan(&i.RootAgentID, &i.ParentAgentID) + return i, err +} + const loadExecutionScope = `-- name: LoadExecutionScope :one SELECT a.root_agent_id FROM agents a WHERE a.project_id=$1 AND a.id=$2 diff --git a/internal/storage/internal/agentexecution/load.go b/internal/storage/internal/agentexecution/load.go index 544db38dd..6d435b6ed 100644 --- a/internal/storage/internal/agentexecution/load.go +++ b/internal/storage/internal/agentexecution/load.go @@ -186,5 +186,5 @@ func executionSnapshot( head.TurnContinuable = selected.TurnContinuable head.IncompleteTools = selected.IncompleteTools head.LogicalReadyAt = selected.LogicalReadyAt - return ExecutionSnapshot{Head: head, View: view, Selection: selected}, nil + return ExecutionSnapshot{Head: head, View: view, Selection: selected, databaseNow: now}, nil } diff --git a/internal/storage/internal/agentexecution/load_count_integration_test.go b/internal/storage/internal/agentexecution/load_count_integration_test.go index 18835602e..c58d2707d 100644 --- a/internal/storage/internal/agentexecution/load_count_integration_test.go +++ b/internal/storage/internal/agentexecution/load_count_integration_test.go @@ -4,6 +4,7 @@ package agentexecution import ( "context" + "strings" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgconn" @@ -14,6 +15,8 @@ type countingLoadTx struct { statements int batches int queries int + facts *pgx.QueuedQuery + factLoads int } func (tx *countingLoadTx) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { @@ -52,6 +55,12 @@ func CountOrdinaryLoad(ctx context.Context, h *Handle) (ExecutionSnapshot, int, func (tx *countingLoadTx) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { tx.statements += batch.Len() + for _, query := range batch.QueuedQueries { + if strings.HasPrefix(query.SQL, "-- name: LoadExecutionFacts ") { + tx.facts = query + tx.factLoads++ + } + } tx.batches++ return tx.Tx.SendBatch(ctx, batch) } @@ -63,3 +72,60 @@ func CountCommit(ctx context.Context, u *Unit) (int, int, error) { err := u.Commit(ctx, "count publication") return tx.statements, tx.batches + tx.queries, err } + +func ReloadExecution(ctx context.Context, h *Handle) (ExecutionSnapshot, error) { + route, err := h.Route() + if err != nil { + return ExecutionSnapshot{}, err + } + var head *ExecutionHead + if h.mutation != nil { + head = &h.mutation.head + } + return loadExecution(ctx, h.unit.DB(), route, head) +} + +func ExplainOrdinaryLoad(ctx context.Context, h *Handle) (string, error) { + tx := &countingLoadTx{Tx: h.unit.tx} + h.unit.tx = tx + _, err := ReloadExecution(ctx, h) + h.unit.tx = tx.Tx + if err != nil { + return "", err + } + if _, err := tx.Tx.Exec(ctx, "SET LOCAL plan_cache_mode=force_generic_plan"); err != nil { + return "", err + } + var plan strings.Builder + for range 10 { + plan.Reset() + rows, err := tx.Tx.Query(ctx, "EXPLAIN (ANALYZE, BUFFERS, TIMING OFF) "+tx.facts.SQL, tx.facts.Arguments...) + if err != nil { + return "", err + } + for rows.Next() { + var line string + if err := rows.Scan(&line); err != nil { + rows.Close() + return "", err + } + plan.WriteString(line) + plan.WriteByte('\n') + } + err = rows.Err() + rows.Close() + if err != nil { + return "", err + } + } + return plan.String(), nil +} + +func CountExecutionTransition(ctx context.Context, u *Unit, run func()) (int, int, error) { + tx := &countingLoadTx{Tx: u.tx} + u.tx = tx + defer func() { u.tx = tx.Tx }() + run() + err := u.Commit(ctx, "count transition") + return tx.statements, tx.factLoads, err +} diff --git a/internal/storage/internal/agentexecution/load_integration_test.go b/internal/storage/internal/agentexecution/load_integration_test.go index 32a3d3570..62de7ad99 100644 --- a/internal/storage/internal/agentexecution/load_integration_test.go +++ b/internal/storage/internal/agentexecution/load_integration_test.go @@ -40,6 +40,16 @@ func TestOrdinaryExecutionLoadStatementBound(t *testing.T) { snapshot, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) require.NoError(t, err) require.Equal(t, 2, count) + _, count, err = agentexecution.CountOrdinaryLoad(t.Context(), h) + require.NoError(t, err) + expected := 0 + if snapshot.Selection.Wait == agentexecution.WaitModelDeadline { + expected = 1 + } + require.Equal(t, expected, count) + plan, err := agentexecution.ExplainOrdinaryLoad(t.Context(), h) + require.NoError(t, err) + t.Log(plan) if snapshot.View.Turn != nil && name != "initial" { require.Empty(t, snapshot.View.Turn.InitialOpening.InputIDs) } diff --git a/internal/storage/internal/agentexecution/locks.go b/internal/storage/internal/agentexecution/locks.go index 194d5a295..2ee876918 100644 --- a/internal/storage/internal/agentexecution/locks.go +++ b/internal/storage/internal/agentexecution/locks.go @@ -121,12 +121,15 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo continue } var err error + var identity dbsqlc.LockAgentInProjectRow if skipLocked { - _, err = q.TryLockAgentInProject(ctx, dbsqlc.TryLockAgentInProjectParams{ + var row dbsqlc.TryLockAgentInProjectRow + row, err = q.TryLockAgentInProject(ctx, dbsqlc.TryLockAgentInProjectParams{ ProjectID: route.ProjectID, ID: route.AgentID, }) + identity = dbsqlc.LockAgentInProjectRow(row) } else { - _, err = q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ + identity, err = q.LockAgentInProject(ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: route.ProjectID, ID: route.AgentID, }) } @@ -139,20 +142,23 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo if err != nil { return false, err } - u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} + u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true, identity: &identity} } if owned, ok := plan.authority.(RuntimeAuthority); ok { h := u.handles[owned.AgentID] if h == nil || owned.RuntimeLockID == uuid.Nil { return false, storeerr.ErrRuntimeLockInactive } - if _, err := q.LockAgentRuntimeLockForOwnedMutation(ctx, dbsqlc.LockAgentRuntimeLockForOwnedMutationParams{ - ProjectID: h.route.ProjectID, AgentID: owned.AgentID, ID: owned.RuntimeLockID, - }); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return false, storeerr.ErrRuntimeLockInactive + if h.lockedRuntimeID != owned.RuntimeLockID { + if _, err := q.LockAgentRuntimeLockForOwnedMutation(ctx, dbsqlc.LockAgentRuntimeLockForOwnedMutationParams{ + ProjectID: h.route.ProjectID, AgentID: owned.AgentID, ID: owned.RuntimeLockID, + }); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return false, storeerr.ErrRuntimeLockInactive + } + return false, err } - return false, err + h.lockedRuntimeID = owned.RuntimeLockID } u.runtimeLocked = true } @@ -207,12 +213,8 @@ func (u *Unit) ClaimAgent(ctx context.Context) (AgentRoute, error) { if err != nil { return AgentRoute{}, err } - plan, err := u.PlanAgents(ctx, []lifecyclelock.AgentRef{{ProjectID: wakeup.ProjectID, AgentID: wakeup.AgentID}}, - IngressAuthority{}) - if err != nil { - return AgentRoute{}, err - } - route := plan.agents[0] + route := AgentRoute{CellID: u.cell.id, ProjectID: wakeup.ProjectID, + AgentID: wakeup.AgentID, RootAgentID: wakeup.RootAgentID} u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} return route, nil } @@ -250,5 +252,36 @@ func (u *Unit) LockAgent( } return dbsqlc.LockAgentInProjectRow{}, err } - return dbsqlc.New(u.DB()).LockAgentInProject(ctx, arg) + h := u.handles[arg.ID] + if h.identity != nil && (h.mutation == nil || !h.mutation.dirty) { + return *h.identity, nil + } + row, err := dbsqlc.New(u.DB()).LockAgentInProject(ctx, arg) + if err == nil { + h.identity = &row + } + return row, err +} + +func (u *Unit) PlanAgentFamily( + ctx context.Context, projectID, agentID uuid.UUID, authority MutationAuthority, +) (MutationPlan, error) { + if err := u.active(); err != nil { + return MutationPlan{}, err + } + row, err := executiondb.New().LoadExecutionFamily(ctx, u.DB(), executiondb.LoadExecutionFamilyParams{ + ProjectID: projectID, ID: agentID, + }) + if errors.Is(err, pgx.ErrNoRows) { + return MutationPlan{}, storeerr.ErrNotFound + } + if err != nil { + return MutationPlan{}, err + } + routes := []AgentRoute{{CellID: u.cell.id, ProjectID: projectID, AgentID: agentID, RootAgentID: row.RootAgentID}} + if row.ParentAgentID != nil { + routes = append(routes, AgentRoute{CellID: u.cell.id, ProjectID: projectID, + AgentID: *row.ParentAgentID, RootAgentID: row.RootAgentID}) + } + return MutationPlan{unit: u, agents: routes, authority: authority}, nil } diff --git a/internal/storage/internal/agentexecution/mutation.go b/internal/storage/internal/agentexecution/mutation.go index 51f40e7a5..83fcc3150 100644 --- a/internal/storage/internal/agentexecution/mutation.go +++ b/internal/storage/internal/agentexecution/mutation.go @@ -6,20 +6,25 @@ import ( "errors" "fmt" - "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" ) type executionMutation struct { - head ExecutionHead - insertHead bool - dirty bool - loaded *ExecutionSnapshot + head ExecutionHead + insertHead bool + dirty bool + selectionChanged bool + loaded *ExecutionSnapshot } func (m *executionMutation) clone() *executionMutation { cloned := *m - cloned.loaded = nil + if m.loaded != nil { + snapshot := cloneExecutionSnapshot(*m.loaded) + cloned.loaded = &snapshot + } if m.head.LogicalReadyAt != nil { ready := *m.head.LogicalReadyAt cloned.head.LogicalReadyAt = &ready @@ -34,11 +39,16 @@ func (h *Handle) executionMutation(ctx context.Context) (*executionMutation, err if h.mutation != nil { return h.mutation, nil } - snapshot, err := h.LoadExecution(ctx) + base, err := executiondb.New().LoadExecutionBase(ctx, h.unit.DB(), executiondb.LoadExecutionBaseParams{ + ProjectID: h.route.ProjectID, AgentID: h.route.AgentID, + }) if err != nil { return nil, err } - m := &executionMutation{head: snapshot.Head, loaded: &snapshot} + if !base.HeadExists { + return nil, fmt.Errorf("load execution head: %w", pgx.ErrNoRows) + } + m := &executionMutation{head: executionHead(base)} h.mutation = m return m, nil } @@ -68,60 +78,52 @@ func (m *executionMutation) publicationParams(route AgentRoute) executiondb.Publ func (u *Unit) publishExecution(ctx context.Context) error { var handles []*Handle var loads []executiondb.LoadExecutionFactsParams + var reloaded []*Handle for _, h := range u.orderedHandles() { if h.mutation == nil || !h.mutation.dirty { continue } handles = append(handles, h) - args := executionFactsParams(h.route, h.mutation.head) - _, captured := u.webhooks[h.route.AgentID] - args.CaptureWebhook = !captured - loads = append(loads, args) + if h.mutation.loaded == nil && h.mutation.selectionChanged { + reloaded = append(reloaded, h) + loads = append(loads, executionFactsParams(h.route, h.mutation.head)) + } } if len(handles) == 0 { return nil } var failure error q := executiondb.New() - batch := q.LoadExecutionFacts(ctx, u.DB(), loads) - batch.Query(func(i int, rows []executiondb.LoadExecutionFactsRow, err error) { - if failure != nil { - return - } - if err != nil { - failure = err - return - } - h := handles[i] - if loads[i].CaptureWebhook { - for _, row := range rows { - if row.Kind != "scope" { - continue - } - target := webhookTarget{orgID: row.WebhookOrgID} - if err := json.Unmarshal(row.WebhookEvents, &target.events); err != nil { - failure = err - return - } - u.webhooks[h.route.AgentID] = target + if len(loads) > 0 { + batch := q.LoadExecutionFacts(ctx, u.DB(), loads) + batch.Query(func(i int, rows []executiondb.LoadExecutionFactsRow, err error) { + if failure != nil { + return } + if err != nil { + failure = err + return + } + h := reloaded[i] + snapshot, err := executionSnapshot(h.route, h.mutation.head, rows) + if err != nil { + failure = err + return + } + h.mutation.head = snapshot.Head + }) + if err := errors.Join(failure, batch.Close()); err != nil { + return err } - snapshot, err := executionSnapshot(h.route, h.mutation.head, rows) - if err != nil { - failure = err - return - } - h.mutation.head = snapshot.Head - }) - if err := errors.Join(failure, batch.Close()); err != nil { - return err } writes := make([]executiondb.PublishExecutionParams, len(handles)) for i, h := range handles { writes[i] = h.mutation.publicationParams(h.route) + _, captured := u.webhooks[h.route.AgentID] + writes[i].CaptureWebhook = !captured } published := q.PublishExecution(ctx, u.DB(), writes) - published.QueryRow(func(i int, id uuid.UUID, err error) { + published.QueryRow(func(i int, row executiondb.PublishExecutionRow, err error) { if failure != nil { return } @@ -129,8 +131,17 @@ func (u *Unit) publishExecution(ctx context.Context) error { failure = err return } - if id != handles[i].route.AgentID { + if row.AgentID != handles[i].route.AgentID { failure = fmt.Errorf("%w: missing execution head", ErrInvalidState) + return + } + if writes[i].CaptureWebhook { + target := webhookTarget{orgID: row.WebhookOrgID} + if err := json.Unmarshal(row.WebhookEvents, &target.events); err != nil { + failure = err + return + } + u.webhooks[row.AgentID] = target } }) return errors.Join(failure, published.Close()) @@ -155,5 +166,30 @@ func executeCommand[T any]( func (m *executionMutation) changed() { m.dirty = true + m.selectionChanged = true m.loaded = nil } + +func (m *executionMutation) selectLoaded() error { + snapshot := m.loaded + snapshot.View.StopSequence = m.head.StopSequence + snapshot.View.MaxNormalInputSequence = m.head.MaxNormalInputSequence + snapshot.View.MaxContextInputSequence = m.head.MaxContextInputSequence + selected, err := SelectNext(snapshot.View, snapshot.databaseNow) + if err != nil { + return err + } + snapshot.Selection = selected + m.head.TurnContinuable = snapshot.Selection.TurnContinuable + m.head.IncompleteTools = snapshot.Selection.IncompleteTools + m.head.LogicalReadyAt = copyPointer(snapshot.Selection.LogicalReadyAt) + snapshot.Head = m.head + m.selectionChanged = false + return nil +} + +func (m *executionMutation) updated(snapshot ExecutionSnapshot) error { + m.dirty = true + m.loaded = &snapshot + return m.selectLoaded() +} diff --git a/internal/storage/internal/agentexecution/publication_integration_test.go b/internal/storage/internal/agentexecution/publication_integration_test.go index cb13317d0..8d703b219 100644 --- a/internal/storage/internal/agentexecution/publication_integration_test.go +++ b/internal/storage/internal/agentexecution/publication_integration_test.go @@ -45,8 +45,12 @@ func TestExecutionCommitRoundTrips(t *testing.T) { } statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) require.NoError(t, err) - require.Equal(t, 2, statements) - require.Equal(t, 2, roundTrips) + expected := 2 + if step == "model" { + expected = 1 + } + require.Equal(t, expected, statements) + require.Equal(t, expected, roundTrips) }) } for _, size := range []int{2, 1000} { diff --git a/internal/storage/internal/agentexecution/queries/load.sql b/internal/storage/internal/agentexecution/queries/load.sql index b7179322b..f50b5ccee 100644 --- a/internal/storage/internal/agentexecution/queries/load.sql +++ b/internal/storage/internal/agentexecution/queries/load.sql @@ -19,15 +19,17 @@ WITH scope AS MATERIALIZED ( AND i.delivery_mode='steering' AND i.state='received') AS steering, EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' AND i.delivery_mode='queued' AND i.state='received') AS queued, - coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event - ORDER BY x.sequence LIMIT 1),0)::bigint AS first_opening_sequence, - coalesce((SELECT x.sequence FROM agent_events x WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event - ORDER BY x.sequence DESC LIMIT 1),0)::bigint AS last_opening_sequence, - coalesce((SELECT x.sequence FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id - WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event AND i.input_kind='content' - ORDER BY x.sequence LIMIT 1),0)::bigint AS first_content_sequence + coalesce(o.first_opening_sequence,0)::bigint AS first_opening_sequence, + coalesce(o.last_opening_sequence,0)::bigint AS last_opening_sequence, + coalesce(o.first_content_sequence,0)::bigint AS first_content_sequence FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=sqlc.narg(turn_id) LEFT JOIN agent_events e ON e.agent_id=a.id AND e.id=t.latest_semantic_event_id + LEFT JOIN LATERAL ( + SELECT min(x.sequence) AS first_opening_sequence,max(x.sequence) AS last_opening_sequence, + min(x.sequence) FILTER (WHERE i.input_kind='content') AS first_content_sequence + FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id + WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event + ) o ON true WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) ), outputs AS MATERIALIZED ( SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at @@ -91,9 +93,9 @@ SELECT 'context'::text AS kind, 'epoch'::timestamptz AS database_now, false AS steering,false AS queued, 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM model_call_contexts c WHERE c.agent_id=sqlc.arg(agent_id) AND c.id IN ( - SELECT sqlc.narg(normal_context_id)::uuid UNION SELECT sqlc.narg(compaction_context_id)::uuid - UNION SELECT model_call_context_id FROM outputs) +FROM model_call_contexts c WHERE c.agent_id=sqlc.arg(agent_id) AND c.id=ANY( + ARRAY[sqlc.narg(normal_context_id)::uuid,sqlc.narg(compaction_context_id)::uuid] + || ARRAY(SELECT model_call_context_id FROM outputs)) UNION ALL SELECT 'output'::text AS kind, o.id AS id, @@ -258,14 +260,8 @@ FROM scope s SELECT f.kind,f.id,f.turn_id,f.context_id,f.event_id,f.sequence,f.event_time,f.operation,f.attempt, f.watermark,f.source_end,f.state,f.recovery,f.retry_at,f.opening_ids,f.opening_sequence, f.has_tools,f.incomplete,f.runnable,f.calls,f.results,f.events,f.stop_reason,f.root_agent_id, - f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence, - coalesce(w.org_id,'00000000-0000-0000-0000-000000000000'::uuid)::uuid AS webhook_org_id, - coalesce(w.events,'[]'::jsonb)::jsonb AS webhook_events -FROM facts f LEFT JOIN LATERAL ( - SELECT p.org_id,c.compiled_definition->'event_webhook'->'events' AS events - FROM agents a JOIN agent_configs c ON c.project_id=a.project_id AND c.id=a.current_config_id - JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id - WHERE a.id=sqlc.arg(agent_id) AND p.deleted_at IS NULL AND o.deleted_at IS NULL - AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' - AND sqlc.arg(capture_webhook)::boolean AND f.kind='scope' -) w ON true ORDER BY f.kind,f.sequence; + f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence +FROM facts f ORDER BY f.kind,f.sequence; + +-- name: ExecutionDatabaseTime :one +SELECT statement_timestamp()::timestamptz; diff --git a/internal/storage/internal/agentexecution/queries/publication.sql b/internal/storage/internal/agentexecution/queries/publication.sql index f321ba741..490008bb2 100644 --- a/internal/storage/internal/agentexecution/queries/publication.sql +++ b/internal/storage/internal/agentexecution/queries/publication.sql @@ -40,4 +40,14 @@ WHERE agent_wakeups.ready_at IS DISTINCT FROM CASE WHEN agent_wakeups.ready_at<=statement_timestamp() AND excluded.ready_at<=statement_timestamp() THEN least(agent_wakeups.ready_at,excluded.ready_at) ELSE excluded.ready_at END RETURNING agent_id -) SELECT agent_id FROM scope; +) SELECT scope.agent_id, + coalesce(w.org_id,'00000000-0000-0000-0000-000000000000'::uuid)::uuid AS webhook_org_id, + coalesce(w.events,'[]'::jsonb)::jsonb AS webhook_events +FROM scope LEFT JOIN LATERAL ( + SELECT p.org_id,c.compiled_definition->'event_webhook'->'events' AS events + FROM agents a JOIN agent_configs c ON c.project_id=a.project_id AND c.id=a.current_config_id + JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id + WHERE a.id=scope.agent_id AND p.deleted_at IS NULL AND o.deleted_at IS NULL + AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' + AND sqlc.arg(capture_webhook)::boolean +) w ON true; diff --git a/internal/storage/internal/agentexecution/queries/state.sql b/internal/storage/internal/agentexecution/queries/state.sql index 0688f8aa4..3d00becc3 100644 --- a/internal/storage/internal/agentexecution/queries/state.sql +++ b/internal/storage/internal/agentexecution/queries/state.sql @@ -11,3 +11,6 @@ SELECT coalesce((SELECT max(e.sequence) FROM agent_events e ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS output_sequence, coalesce((SELECT t.id FROM agent_turns t WHERE t.agent_id=$2 ORDER BY t.turn_sequence DESC LIMIT 1), '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS turn_id; + +-- name: LoadExecutionFamily :one +SELECT root_agent_id,parent_agent_id FROM agents WHERE project_id=$1 AND id=$2; diff --git a/internal/storage/internal/agentexecution/schema_integration_test.go b/internal/storage/internal/agentexecution/schema_integration_test.go index 456a00637..94227c35f 100644 --- a/internal/storage/internal/agentexecution/schema_integration_test.go +++ b/internal/storage/internal/agentexecution/schema_integration_test.go @@ -133,6 +133,16 @@ created_at,updated_at) ) require.Equal(t, tc.valid, valid) } + for _, array := range []string{ + "'[0:0]={11111111-1111-4111-8111-111111111111}'", + "'{{11111111-1111-4111-8111-111111111111}}'", "'{NULL}'", + } { + var valid bool + require.NoError(t, f.pool.QueryRow(t.Context(), + "SELECT execution_opening_is_valid("+array+"::uuid[],1)").Scan(&valid)) + require.False(t, valid) + } + } func TestExecutionReconstructionIgnoresCorruptHead(t *testing.T) { @@ -184,6 +194,8 @@ SELECT normal_context_id IS NULL FROM agent_execution_state WHERE agent_id=$1 require.Equal(t, agentexecution.WorkModel, rebuilt.Selection.Work) _, err = u.DB().Exec(t.Context(), `DELETE FROM agent_execution_state WHERE agent_id=$1`, agent) require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "corrupt head")) + _, h = f.handle(t, agent) _, err = h.LoadExecution(t.Context()) require.ErrorIs(t, err, pgx.ErrNoRows) rebuilt, err = h.ReconstructExecution(t.Context()) diff --git a/internal/storage/internal/agentexecution/selection_load.go b/internal/storage/internal/agentexecution/selection_load.go index a05a6482d..89351d36f 100644 --- a/internal/storage/internal/agentexecution/selection_load.go +++ b/internal/storage/internal/agentexecution/selection_load.go @@ -28,9 +28,10 @@ type ExecutionHead struct { } type ExecutionSnapshot struct { - Head ExecutionHead - View ExecutionView - Selection Selection + Head ExecutionHead + View ExecutionView + Selection Selection + databaseNow time.Time } type executionLoader struct { @@ -68,13 +69,29 @@ func (h *Handle) LoadExecution(ctx context.Context) (ExecutionSnapshot, error) { m := h.mutation if m != nil { head = &m.head - if m.loaded != nil && m.loaded.Selection.Wait != WaitModelDeadline { + if m.loaded != nil { + if m.loaded.Selection.Wait == WaitModelDeadline { + now, err := executiondb.New().ExecutionDatabaseTime(ctx, h.unit.DB()) + if err != nil { + return ExecutionSnapshot{}, err + } + m.loaded.databaseNow = now + if err := m.selectLoaded(); err != nil { + return ExecutionSnapshot{}, err + } + } return cloneExecutionSnapshot(*m.loaded), nil } } snapshot, err := loadExecution(ctx, h.unit.DB(), route, head) - if err == nil && m != nil { + if err == nil { + if m == nil { + m = &executionMutation{} + h.mutation = m + } + m.head = snapshot.Head m.loaded = &snapshot + m.selectionChanged = false } return cloneExecutionSnapshot(snapshot), err } diff --git a/internal/storage/internal/agentexecution/transition_count_integration_test.go b/internal/storage/internal/agentexecution/transition_count_integration_test.go new file mode 100644 index 000000000..a6ddd2d2a --- /dev/null +++ b/internal/storage/internal/agentexecution/transition_count_integration_test.go @@ -0,0 +1,78 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "testing" + + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/require" +) + +func TestExecutionTransitionStatementBounds(t *testing.T) { + f, a, lease := commandFixture(t) + measured := func(name string, run func(*agentexecution.Handle)) { + u, h := f.handle(t, a) + statements, facts, err := agentexecution.CountExecutionTransition(t.Context(), u, func() { run(h) }) + require.NoError(t, err) + t.Logf("%s statements=%d facts=%d", name, statements, facts) + limits := map[string]int{"receive": 6, "prepare_initial": 13, "accept_tools": 12, + "authorize": 6, "run_tool": 7, "complete_prepare": 16, "accept_end": 11} + require.LessOrEqual(t, statements, limits[name], name) + maximumFacts := 1 + if name == "authorize" { + maximumFacts = 0 + } + require.LessOrEqual(t, facts, maximumFacts, name) + } + const cycles = 2 + for range cycles { + var prepared agentexecution.PreparedModel + var output agentexecution.AcceptedOutput + measured("receive", func(h *agentexecution.Handle) { receiveCommand(t, h, "queued", "input") }) + measured("prepare_initial", func(h *agentexecution.Handle) { + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared = prepareCommand(t, h, lease) + }) + measured("accept_tools", func(h *agentexecution.Handle) { + response := commandResponse(modelenvelope.StopReasonToolUse) + response.Normalized.Content = []modelenvelope.ResponsePart{{Type: "tool_call", ProviderCallID: "one", ToolName: "tool", ToolInput: json.RawMessage(`{}`)}} + var err error + output, err = h.AcceptOutput(t.Context(), agentexecution.AcceptOutputInput{ + ContextID: prepared.Context.ID, RuntimeLockID: lease, Response: response, + Tools: []agentexecution.ToolProposal{{ProviderCallID: "one", Type: "built_in"}}, + }) + require.NoError(t, err) + _, err = h.LoadExecution(t.Context()) + require.NoError(t, err) + }) + measured("authorize", func(h *agentexecution.Handle) { + _, err := h.AuthorizeTool(t.Context(), agentexecution.ToolRef{ID: output.Tools[0].ID, RuntimeLockID: lease}) + require.NoError(t, err) + }) + measured("run_tool", func(h *agentexecution.Handle) { + _, err := h.RunTool(t.Context(), agentexecution.ToolRef{ID: output.Tools[0].ID, RuntimeLockID: lease}) + require.NoError(t, err) + }) + measured("complete_prepare", func(h *agentexecution.Handle) { + _, err := h.CompleteTool(t.Context(), agentexecution.ToolCompletion{ + ToolRef: agentexecution.ToolRef{ID: output.Tools[0].ID, RuntimeLockID: lease}, + Outcome: "succeeded", Content: []agentexecution.Content{{Kind: "text", Text: "done"}}, + }) + require.NoError(t, err) + prepared = prepareCommand(t, h, lease) + }) + measured("accept_end", func(h *agentexecution.Handle) { + _, err := h.AcceptOutput(t.Context(), agentexecution.AcceptOutputInput{ + ContextID: prepared.Context.ID, RuntimeLockID: lease, + Response: commandResponse(modelenvelope.StopReasonEndTurn), + }) + require.NoError(t, err) + _, err = h.LoadExecution(t.Context()) + require.NoError(t, err) + }) + } +} diff --git a/internal/storage/internal/dbsqlc/agent_inputs.sql.go b/internal/storage/internal/dbsqlc/agent_inputs.sql.go index fd14a3a9b..d4f23763f 100644 --- a/internal/storage/internal/dbsqlc/agent_inputs.sql.go +++ b/internal/storage/internal/dbsqlc/agent_inputs.sql.go @@ -15,7 +15,7 @@ import ( const claimNextAgentWakeup = `-- name: ClaimNextAgentWakeup :one WITH locked_agent AS MATERIALIZED ( - SELECT agent.id AS agent_id, agent.project_id + SELECT agent.id AS agent_id, agent.project_id, agent.root_agent_id FROM agent_wakeups wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.state <> 'archived' @@ -30,7 +30,7 @@ WITH locked_agent AS MATERIALIZED ( LIMIT 1 ), locked_wake AS MATERIALIZED ( - SELECT wake.agent_id, agent.project_id + SELECT wake.agent_id, agent.project_id, agent.root_agent_id FROM agent_wakeups wake JOIN locked_agent agent ON agent.agent_id = wake.agent_id WHERE NOT EXISTS ( @@ -40,13 +40,14 @@ locked_wake AS MATERIALIZED ( ) FOR UPDATE OF wake ) -SELECT agent_id, project_id +SELECT agent_id, project_id, root_agent_id FROM locked_wake ` type ClaimNextAgentWakeupRow struct { - AgentID uuid.UUID - ProjectID uuid.UUID + AgentID uuid.UUID + ProjectID uuid.UUID + RootAgentID uuid.UUID } // The query walks wakeups in global ready order, but locks agents before wake @@ -56,7 +57,7 @@ type ClaimNextAgentWakeupRow struct { func (q *Queries) ClaimNextAgentWakeup(ctx context.Context) (ClaimNextAgentWakeupRow, error) { row := q.db.QueryRow(ctx, claimNextAgentWakeup) var i ClaimNextAgentWakeupRow - err := row.Scan(&i.AgentID, &i.ProjectID) + err := row.Scan(&i.AgentID, &i.ProjectID, &i.RootAgentID) return i, err } diff --git a/internal/storage/internal/dbsqlc/agent_runtime.sql.go b/internal/storage/internal/dbsqlc/agent_runtime.sql.go index 097d48e8a..d5673d61d 100644 --- a/internal/storage/internal/dbsqlc/agent_runtime.sql.go +++ b/internal/storage/internal/dbsqlc/agent_runtime.sql.go @@ -722,7 +722,7 @@ func (q *Queries) LockAgentLaunchIdempotencyKey(ctx context.Context, arg LockAge } const tryLockAgentInProject = `-- name: TryLockAgentInProject :one -SELECT id, org_id +SELECT id, org_id, state FROM agents WHERE project_id = $1 AND id = $2 FOR UPDATE SKIP LOCKED @@ -736,11 +736,12 @@ type TryLockAgentInProjectParams struct { type TryLockAgentInProjectRow struct { ID uuid.UUID OrgID uuid.UUID + State string } func (q *Queries) TryLockAgentInProject(ctx context.Context, arg TryLockAgentInProjectParams) (TryLockAgentInProjectRow, error) { row := q.db.QueryRow(ctx, tryLockAgentInProject, arg.ProjectID, arg.ID) var i TryLockAgentInProjectRow - err := row.Scan(&i.ID, &i.OrgID) + err := row.Scan(&i.ID, &i.OrgID, &i.State) return i, err } diff --git a/internal/storage/internal/dbsqlc/subagents.sql.go b/internal/storage/internal/dbsqlc/subagents.sql.go index 5630a5fef..67a515e20 100644 --- a/internal/storage/internal/dbsqlc/subagents.sql.go +++ b/internal/storage/internal/dbsqlc/subagents.sql.go @@ -130,25 +130,6 @@ func (q *Queries) CountAgentAncestors(ctx context.Context, arg CountAgentAncesto return column_1, err } -const getAgentParentID = `-- name: GetAgentParentID :one -SELECT parent_agent_id -FROM agents -WHERE project_id = $1 - AND id = $2 -` - -type GetAgentParentIDParams struct { - ProjectID uuid.UUID - ID uuid.UUID -} - -func (q *Queries) GetAgentParentID(ctx context.Context, arg GetAgentParentIDParams) (*uuid.UUID, error) { - row := q.db.QueryRow(ctx, getAgentParentID, arg.ProjectID, arg.ID) - var parent_agent_id *uuid.UUID - err := row.Scan(&parent_agent_id) - return parent_agent_id, err -} - const listActiveChildAgentIDs = `-- name: ListActiveChildAgentIDs :many SELECT agent.id FROM agents agent diff --git a/internal/storage/queries/agent_inputs.sql b/internal/storage/queries/agent_inputs.sql index 8cdac3ce6..5babb0d4e 100644 --- a/internal/storage/queries/agent_inputs.sql +++ b/internal/storage/queries/agent_inputs.sql @@ -4,7 +4,7 @@ -- longer claimable after the runtime-lock recheck, this returns no rows and the -- worker retries on its next poll; no work is lost. WITH locked_agent AS MATERIALIZED ( - SELECT agent.id AS agent_id, agent.project_id + SELECT agent.id AS agent_id, agent.project_id, agent.root_agent_id FROM agent_wakeups wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.state <> 'archived' @@ -19,7 +19,7 @@ WITH locked_agent AS MATERIALIZED ( LIMIT 1 ), locked_wake AS MATERIALIZED ( - SELECT wake.agent_id, agent.project_id + SELECT wake.agent_id, agent.project_id, agent.root_agent_id FROM agent_wakeups wake JOIN locked_agent agent ON agent.agent_id = wake.agent_id WHERE NOT EXISTS ( @@ -29,7 +29,7 @@ locked_wake AS MATERIALIZED ( ) FOR UPDATE OF wake ) -SELECT agent_id, project_id +SELECT agent_id, project_id, root_agent_id FROM locked_wake; -- name: ListQueuedBacklogInputs :many diff --git a/internal/storage/queries/agent_runtime.sql b/internal/storage/queries/agent_runtime.sql index 4132b5b45..4b12afd88 100644 --- a/internal/storage/queries/agent_runtime.sql +++ b/internal/storage/queries/agent_runtime.sql @@ -251,7 +251,7 @@ WHERE project_id = $1 AND id = $2 FOR UPDATE; -- name: TryLockAgentInProject :one -SELECT id, org_id +SELECT id, org_id, state FROM agents WHERE project_id = $1 AND id = $2 FOR UPDATE SKIP LOCKED; diff --git a/internal/storage/queries/subagents.sql b/internal/storage/queries/subagents.sql index 04f6f8769..0c51cf8e4 100644 --- a/internal/storage/queries/subagents.sql +++ b/internal/storage/queries/subagents.sql @@ -94,12 +94,6 @@ WHERE agent.project_id = sqlc.arg(project_id) AND agent.parent_agent_id = sqlc.arg(parent_agent_id) AND agent.state = 'active'; --- name: GetAgentParentID :one -SELECT parent_agent_id -FROM agents -WHERE project_id = sqlc.arg(project_id) - AND id = sqlc.arg(id); - -- name: ListActiveChildAgentIDs :many SELECT agent.id FROM agents agent diff --git a/migrations/000052_agent_execution_state.sql b/migrations/000052_agent_execution_state.sql index 4d2b9ebdd..0c566cf39 100644 --- a/migrations/000052_agent_execution_state.sql +++ b/migrations/000052_agent_execution_state.sql @@ -117,12 +117,19 @@ ALTER TABLE context_checkpoints ENABLE TRIGGER context_checkpoints_immutable; -- +goose StatementBegin CREATE FUNCTION execution_opening_is_valid(ids uuid[], first_sequence bigint) -RETURNS boolean LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$ - SELECT coalesce(ids IS NOT NULL AND array_ndims(ids) IS NOT DISTINCT FROM CASE WHEN cardinality(ids)=0 THEN NULL ELSE 1 END - AND (cardinality(ids)=0 AND first_sequence IS NULL OR cardinality(ids)>0 AND first_sequence>0) - AND NOT EXISTS(SELECT 1 FROM unnest(ids) id WHERE id IS NULL OR id='00000000-0000-0000-0000-000000000000'::uuid) - AND (cardinality(ids)=0 OR array_lower(ids,1)=1) - AND cardinality(ids)=(SELECT count(DISTINCT id) FROM unnest(ids) id), false) +RETURNS boolean LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE AS $$ +BEGIN + IF ids IS NULL THEN RETURN false; END IF; + IF cardinality(ids)=0 THEN RETURN first_sequence IS NULL; END IF; + IF first_sequence IS NULL OR first_sequence<=0 OR array_ndims(ids)<>1 OR array_lower(ids,1)<>1 THEN + RETURN false; + END IF; + IF cardinality(ids)=1 THEN + RETURN ids[1] IS NOT NULL AND ids[1]<>'00000000-0000-0000-0000-000000000000'::uuid; + END IF; + RETURN NOT EXISTS(SELECT 1 FROM unnest(ids) id WHERE id IS NULL OR id='00000000-0000-0000-0000-000000000000'::uuid) + AND cardinality(ids)=(SELECT count(DISTINCT id) FROM unnest(ids) id); +END; $$; -- +goose StatementEnd From 170deada64b20ffbafbd0319396c9fdd0e8fe1e7 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:01:07 -0700 Subject: [PATCH 04/19] fix: address agent execution review findings - Check the continuation budget only before preparing the next step, so a prepared step always runs. - Accept model output before optional advancement. Advance inside a savepoint; on failure, commit acceptance alone and surface the cause to the kernel, which logs it. - Archive cancels in-flight work before marking the agent archived, and facts load independently of lifecycle state. - Derive organization-deletion and grant-revocation lock sets from the rows they clean up. A unit that needs an agent outside its plan rolls back and restarts with the expanded sorted plan, and units can only be opened inside that retry boundary. - Lock only the target agent for ordinary transitions so siblings do not serialize on their parent. - Enforce opening-lineage membership and order with triggers. - Replace the discriminated fact UNION with typed queries in one batch. - Remove the unused answered boundary, cell identity and authority tags. - Suppress webhooks for deleted projects and organizations and document it. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/events/webhooks.mdx | 2 + .../harness/kernel/agent_executor_context.go | 10 + .../continuation_failures_integration_test.go | 196 +++++ internal/harness/worker/continuation_test.go | 4 +- internal/harness/worker/worker.go | 2 +- ...default_reconciliation_integration_test.go | 2 +- .../agent_archive_integration_test.go | 21 +- ...nt_config_change_gates_integration_test.go | 3 +- .../agent_config_changes_store.go | 11 +- .../agent_content_inputs_store.go | 58 +- .../agent_fixture_integration.go | 169 +++-- ...agent_input_admission_test_helpers_test.go | 2 +- .../agent_inputs_backlog_store.go | 41 +- .../executionstore/agent_inputs_store.go | 87 ++- .../agent_interactions_store.go | 91 +-- .../executionstore/agent_launch_inbox.go | 10 +- .../executionstore/agent_launch_store.go | 9 +- .../agent_runtime_cancellation_store.go | 42 +- .../agent_runtime_lock_reaping_store.go | 32 +- .../agent_runtime_locks_store.go | 126 ++-- .../executionstore/agent_runtime_store.go | 12 +- .../executionstore/agent_work_store.go | 60 +- .../context_checkpoints_store.go | 71 +- .../context_checkpoints_terminal_failure.go | 47 +- .../cron_trigger_agent_inputs_store.go | 47 +- .../executionstore/execution_mutation.go | 3 +- ...cution_unit_lock_order_integration_test.go | 24 +- ...xecution_unit_webhooks_integration_test.go | 27 +- .../executionstore/export_integration_test.go | 151 ++-- .../executionstore/inbox_input_store.go | 11 +- .../integration_input_origin.go | 2 +- ...tegration_interactions_integration_test.go | 7 +- .../integration_interactions_store.go | 232 +++--- ...tion_selection_queries_integration_test.go | 3 +- ...vent_authority_helpers_integration_test.go | 3 +- .../lifecycle_residual_integration_test.go | 96 +++ .../machine_daemon_bootstrap_store.go | 108 +-- .../machine_daemon_reconciliation_store.go | 4 +- .../machine_daemon_runtime_expiry_store.go | 94 +-- .../machine_daemon_runtime_store.go | 679 +++++++++--------- .../executionstore/machine_daemon_store.go | 25 +- .../executionstore/machine_lifecycle_store.go | 462 ++++++------ .../machine_lifecycle_terminal_work.go | 26 +- .../executionstore/machine_pools_store.go | 22 +- .../executionstore/machine_wake_store.go | 10 +- .../executionstore/model_contexts_claims.go | 102 +-- .../model_contexts_compaction.go | 183 ++--- .../model_contexts_lifecycle.go | 62 +- .../model_contexts_preparation.go | 37 +- ...l_contexts_preparation_integration_test.go | 3 +- .../model_contexts_replay_integration_test.go | 18 +- .../model_outputs_terminal_error.go | 89 +-- .../executionstore/model_outputs_workflow.go | 88 ++- .../executionstore/model_work_transition.go | 64 +- .../model_work_transition_integration_test.go | 93 ++- .../process_action_completion_store.go | 60 +- ...s_cancellation_archive_integration_test.go | 4 +- .../process_daemon_completion_store.go | 472 ++++++------ .../process_tool_unreachable_store.go | 492 +++++++------ .../storage/executionstore/processes_store.go | 94 +-- .../project_machine_pool_grants_store.go | 6 +- .../provider_runtime_reconciliation_store.go | 257 ++++--- internal/storage/executionstore/store.go | 1 - .../storage/executionstore/subagents_store.go | 10 +- .../executionstore/tool_call_execution.go | 11 +- .../executionstore/tool_calls_store.go | 197 ++--- .../integrationstore/export_integration.go | 6 +- .../integrationstore/integration_lifecycle.go | 5 +- ...ation_target_lifecycle_integration_test.go | 2 +- internal/storage/integrationstore/store.go | 2 +- .../storage/internal/agentexecution/cell.go | 35 +- .../agentexecution/commands_agents.go | 5 +- .../agentexecution/commands_cancel.go | 10 - .../commands_family_integration_test.go | 51 +- .../commands_integration_test.go | 44 +- .../agentexecution/commands_lifecycle.go | 15 +- .../agentexecution/commands_models.go | 2 +- .../agentexecution/commands_outputs.go | 1 - .../agentexecution/commands_runtime.go | 18 +- .../commands_runtime_integration_test.go | 15 +- .../internal/agentexecution/commands_tools.go | 2 +- .../internal/agentexecution/facts_batch.go | 121 ++++ .../fixture_integration_test.go | 6 +- .../storage/internal/agentexecution/handle.go | 7 +- .../internal/executiondb/batch.go | 665 +++++++++-------- .../internal/executiondb/load.sql.go | 3 - .../internal/executiondb/state.sql.go | 9 +- .../linearity_integration_test.go | 11 +- .../storage/internal/agentexecution/load.go | 188 ++--- .../load_count_integration_test.go | 15 +- .../agentexecution/load_integration_test.go | 9 +- .../storage/internal/agentexecution/locks.go | 75 +- .../internal/agentexecution/mutation.go | 46 +- .../publication_integration_test.go | 18 +- .../internal/agentexecution/queries/load.sql | 268 ++----- .../agentexecution/queries/publication.sql | 10 +- .../internal/agentexecution/queries/state.sql | 2 - .../reconstruction_integration_test.go | 8 +- .../storage/internal/agentexecution/repair.go | 7 +- .../storage/internal/agentexecution/retry.go | 84 +++ .../storage/internal/agentexecution/route.go | 3 - .../agentexecution/schema_integration_test.go | 70 ++ .../internal/agentexecution/selection_load.go | 7 +- .../agentexecution/testing_integration.go | 9 + .../transition_count_integration_test.go | 4 +- .../storage/internal/agentexecution/unit.go | 16 +- .../agentexecution/unit_integration_test.go | 228 +++++- .../internal/agentexecution/webhooks.go | 7 +- .../internal/dbsqlc/agent_inputs.sql.go | 20 + .../storage/internal/dbsqlc/processes.sql.go | 66 -- .../orglifecycle/export_integration.go | 13 +- .../storage/orglifecycle/orgs_projects.go | 50 +- internal/storage/orglifecycle/service.go | 2 +- internal/storage/queries/agent_inputs.sql | 8 + internal/storage/queries/processes.sql | 33 - migrations/000052_agent_execution_state.sql | 35 +- 116 files changed, 4446 insertions(+), 3377 deletions(-) create mode 100644 internal/harness/worker/continuation_failures_integration_test.go create mode 100644 internal/storage/executionstore/lifecycle_residual_integration_test.go create mode 100644 internal/storage/internal/agentexecution/facts_batch.go create mode 100644 internal/storage/internal/agentexecution/retry.go create mode 100644 internal/storage/internal/agentexecution/testing_integration.go diff --git a/docs/events/webhooks.mdx b/docs/events/webhooks.mdx index 8e3ec548c..efdfde679 100644 --- a/docs/events/webhooks.mdx +++ b/docs/events/webhooks.mdx @@ -6,6 +6,8 @@ icon: "webhook" Configure an [event webhook](/agents/configuration#event-webhook) to receive JSON POST requests with `event` and `data`. The event name determines the data schema. These are the same data objects used by the [event stream](/events/streaming). +Webhook deliveries are suppressed for deleted projects and organizations, including events produced by their deletion. + Every request includes `Webhook-Id` (stable across retries) and `Webhook-Timestamp` (Unix seconds for that attempt). Signed requests also include `Webhook-Signature`. Verify the signature against the raw request body before parsing it, using [Standard Webhooks](https://github.com/standard-webhooks/standard-webhooks/blob/main/spec/standard-webhooks.md). ## Receive and verify diff --git a/internal/harness/kernel/agent_executor_context.go b/internal/harness/kernel/agent_executor_context.go index 39b9ad912..6d9aac492 100644 --- a/internal/harness/kernel/agent_executor_context.go +++ b/internal/harness/kernel/agent_executor_context.go @@ -3,6 +3,7 @@ package kernel import ( "context" "errors" + "log/slog" "time" "github.com/google/uuid" @@ -360,6 +361,15 @@ func (e AgentExecutor) executeModelStepWithAdvance( AllowModelWork: advance.AllowModelWork, PrepareModel: true, }, ) + if transition.Fallback != nil { + log := e.StreamLog + if log == nil { + log = slog.Default() + } + log.WarnContext(ctx, "model output advancement fell back", + "agent_id", input.AgentID, "model_call_context_id", claim.Context.ID, + "runtime_lock_id", input.RuntimeLockID, "error", transition.Fallback) + } if err != nil { return modelStep{}, err } diff --git a/internal/harness/worker/continuation_failures_integration_test.go b/internal/harness/worker/continuation_failures_integration_test.go new file mode 100644 index 000000000..df5e5cceb --- /dev/null +++ b/internal/harness/worker/continuation_failures_integration_test.go @@ -0,0 +1,196 @@ +//go:build integration + +package worker + +import ( + "bytes" + "context" + "log/slog" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgconn" + "github.com/omnara-ai/omnara/internal/harness/kernel" + "github.com/omnara-ai/omnara/internal/model" + "github.com/omnara-ai/omnara/internal/storage" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/stretchr/testify/require" +) + +type delayedOwnedStore struct { + Store + after func() +} + +func (s delayedOwnedStore) AdvanceOwnedAgentWork(ctx context.Context, + input executionstore.AdvanceOwnedAgentWorkInput) (executionstore.ClaimedAgentWork, + bool, + error) { + work, continued, err := s.Store.AdvanceOwnedAgentWork(ctx, input) + if continued && work.Model.Prepared != nil && s.after != nil { + s.after() + } + return work, continued, err +} + +type delayedFusedExecutor struct { + kernel.AgentExecutor + before, after func() + observe func(executionstore.OwnedAgentWorkTransition) +} + +func (e delayedFusedExecutor) ExecuteModelWorkAndAdvance(ctx context.Context, + input kernel.ModelWorkExecution, + options kernel.ModelWorkAdvanceOptions) (executionstore.OwnedAgentWorkTransition, + bool, + error) { + if e.before != nil { + e.before() + } + work, advanced, err := e.AgentExecutor.ExecuteModelWorkAndAdvance(ctx, input, options) + if e.observe != nil { + e.observe(work) + } + if work.Continued && work.Work.Model.Prepared != nil && e.after != nil { + e.after() + } + return work, advanced, err +} + +func TestWorkerContinuationDurationNeverAbandonsPreparedContext(t *testing.T) { + for _, path := range []string{"owned", "fused"} { + for _, boundary := range []string{"before", "during", "after"} { + t.Run(path+"/"+boundary, func(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agent, user := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + createWorkerInput(t, ctx, store, agent, user, "continue", time.Now()) + first := continuationFinalResponse(1) + first.StopReason = model.StopReasonMaxTokens + executor := newObservedContinuationExecutor(store, first, continuationFinalResponse(2)) + pause := func() { <-time.NewTimer(1100 * time.Millisecond).C } + var workerStore Store = store.Execution() + var workerExecutor AgentWorkExecutor = executor + if path == "owned" { + if boundary == "before" { + executor.afterModel = func(context.Context) error { pause(); return nil } + } + if boundary == "after" { + workerStore = delayedOwnedStore{Store: workerStore, after: pause} + } + } else { + fused := delayedFusedExecutor{AgentExecutor: executor.AgentExecutor} + if boundary == "before" { + fused.before = pause + } + if boundary == "after" { + fused.after = pause + } + workerExecutor = fused + } + if boundary == "during" { + _, err := pool.Exec(ctx, `CREATE FUNCTION delay_successor() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN IF EXISTS(SELECT 1 FROM model_outputs WHERE agent_id=NEW.agent_id) THEN + PERFORM pg_sleep(1.1); END IF; RETURN NEW; END $$; + CREATE TRIGGER delay_successor BEFORE INSERT ON model_call_contexts + FOR EACH ROW EXECUTE FUNCTION delay_successor()`) + require.NoError(t, err) + } + worker := NewWorker(workerStore, + workerExecutor, + Options{Continuation: &ContinuationOptions{MaxModelStarts: 2, + MaxDuration: time.Second}}) + requireWorkerClaim(t, ctx, worker, "duration boundary") + var contexts, outputs, failed int + require.NoError(t, + pool.QueryRow(ctx, + `SELECT count(*),count(*) FILTER(WHERE state='failed') FROM model_call_contexts WHERE agent_id=$1`, + agent).Scan(&contexts, + &failed)) + require.Zero(t, failed) + require.NoError(t, pool.QueryRow(ctx, `SELECT count(*) FROM model_outputs WHERE agent_id=$1`, agent).Scan(&outputs)) + require.Equal(t, contexts, outputs) + expected := 2 + if boundary == "before" { + expected = 1 + } + require.Equal(t, expected, outputs) + if boundary == "before" { + requireWorkerClaim(t, ctx, worker, "resume yielded work") + } + require.NoError(t, + pool.QueryRow(ctx, + `SELECT count(*),count(*) FILTER(WHERE state='failed') FROM model_call_contexts WHERE agent_id=$1`, + agent).Scan(&contexts, + &failed)) + require.Equal(t, 2, contexts) + require.Zero(t, failed) + waitNoWorkerWakeup(t, ctx, pool, agent) + }) + } + } +} + +func TestWorkerAcceptsResponseWhenSuccessorPreparationFails(t *testing.T) { + for _, deferred := range []bool{false, true} { + name := "immediate" + if deferred { + name = "deferred" + } + t.Run(name, func(t *testing.T) { + ctx := t.Context() + pool := openWorkerIntegrationDB(t, ctx) + store := storage.NewStore(pool, storage.WithSecretKeyWrapper(workerTestKeyWrapper(t))) + agent, user := createWorkerAgentWithTools(t, ctx, store, time.Now(), "list_agents") + createWorkerInput(t, ctx, store, agent, user, "continue", time.Now()) + first := continuationFinalResponse(1) + first.StopReason = model.StopReasonMaxTokens + executor := newObservedContinuationExecutor(store, first, continuationFinalResponse(2)) + _, err := pool.Exec(ctx, `CREATE FUNCTION reject_successor() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN IF EXISTS(SELECT 1 FROM model_outputs WHERE agent_id=NEW.agent_id) THEN + RAISE EXCEPTION 'reject successor'; END IF; RETURN NEW; END $$`) + require.NoError(t, err) + trigger := `CREATE TRIGGER reject_successor BEFORE INSERT ON model_call_contexts + FOR EACH ROW EXECUTE FUNCTION reject_successor()` + if deferred { + trigger = `CREATE CONSTRAINT TRIGGER reject_successor AFTER INSERT ON model_call_contexts + DEFERRABLE INITIALLY DEFERRED + FOR EACH ROW EXECUTE FUNCTION reject_successor()` + } + _, err = pool.Exec(ctx, trigger) + require.NoError(t, err) + var logged bytes.Buffer + executor.StreamLog = slog.New(slog.NewJSONHandler(&logged, nil)) + var transition executionstore.OwnedAgentWorkTransition + observed := delayedFusedExecutor{AgentExecutor: executor.AgentExecutor, + observe: func(result executionstore.OwnedAgentWorkTransition) { transition = result }} + worker := NewWorker(store.Execution(), observed, Options{}) + worked, runErr := worker.RunOnce(ctx) + require.True(t, worked) + var outputs, failed int + require.NoError(t, pool.QueryRow(ctx, `SELECT count(*) FROM model_outputs WHERE agent_id=$1`, agent).Scan(&outputs)) + require.Equal(t, 1, outputs) + require.NoError(t, runErr) + var cause *pgconn.PgError + require.ErrorAs(t, transition.Fallback, &cause) + require.Equal(t, "reject successor", cause.Message) + require.Contains(t, logged.String(), `"level":"WARN"`) + require.Contains(t, logged.String(), `"agent_id":"`+agent.String()+`"`) + require.Contains(t, logged.String(), `"model_call_context_id":`) + require.Contains(t, logged.String(), `"runtime_lock_id":`) + require.Contains(t, logged.String(), "SQLSTATE P0001") + require.NoError(t, + pool.QueryRow(ctx, + `SELECT count(*) FROM model_call_contexts WHERE agent_id=$1 AND state='failed'`, + agent).Scan(&failed)) + require.Zero(t, failed) + _, err = pool.Exec(ctx, `DROP TRIGGER reject_successor ON model_call_contexts`) + require.NoError(t, err) + requireWorkerClaim(t, ctx, worker, "resume after rejected successor") + require.NoError(t, pool.QueryRow(ctx, `SELECT count(*) FROM model_outputs WHERE agent_id=$1`, agent).Scan(&outputs)) + require.Equal(t, 2, outputs) + waitNoWorkerWakeup(t, ctx, pool, agent) + }) + } +} diff --git a/internal/harness/worker/continuation_test.go b/internal/harness/worker/continuation_test.go index d2fbc90db..1bae1a6ba 100644 --- a/internal/harness/worker/continuation_test.go +++ b/internal/harness/worker/continuation_test.go @@ -218,7 +218,7 @@ func TestWorkerOwnedContinuationKeepsRenewalRunning(t *testing.T) { } } -func TestWorkerOwnedContinuationCountsTimeSpentAdvancing(t *testing.T) { +func TestWorkerOwnedContinuationExecutesWorkAfterSoftDeadline(t *testing.T) { store := newContinuationStore(executionstore.AgentWorkModel) store.advanceDelay = 110 * time.Millisecond executor := &continuationExecutor{} @@ -229,7 +229,7 @@ func TestWorkerOwnedContinuationCountsTimeSpentAdvancing(t *testing.T) { if err != nil || !worked { t.Fatalf("worked=%v err=%v", worked, err) } - if len(executor.kinds) != 1 || len(store.advances) != 1 || len(store.released) != 1 { + if len(executor.kinds) != 2 || len(store.advances) != 1 || len(store.released) != 1 { t.Fatalf("steps=%v advances=%d releases=%d", executor.kinds, len(store.advances), len(store.released)) } } diff --git a/internal/harness/worker/worker.go b/internal/harness/worker/worker.go index d21eecd72..c334d97a4 100644 --- a/internal/harness/worker/worker.go +++ b/internal/harness/worker/worker.go @@ -377,7 +377,7 @@ func (w *Worker) executeClaimedWork( } } err = turnCtx.Err() - if err != nil || !w.continuation.canAdvance(time.Since(startedAt)) { + if err != nil { break } if claim.Kind == executionstore.AgentWorkModel { diff --git a/internal/storage/default_reconciliation_integration_test.go b/internal/storage/default_reconciliation_integration_test.go index 4412a5663..4dd421f89 100644 --- a/internal/storage/default_reconciliation_integration_test.go +++ b/internal/storage/default_reconciliation_integration_test.go @@ -1011,7 +1011,7 @@ INSERT INTO machines( }) integrationdb.WaitForNamedLockWaiters(t, ctx, pool, "LockMachineForLifecycle", 1) deleteDone := integrationdb.RunAsyncError(func() error { - deleteTx, deleteErr := agentexecution.NewCell("primary", pool, nil, nil).Begin(ctx) + deleteTx, deleteErr := agentexecution.BeginTestUnit(ctx, agentexecution.NewCell(pool, nil, nil)) if deleteErr != nil { return deleteErr } diff --git a/internal/storage/executionstore/agent_archive_integration_test.go b/internal/storage/executionstore/agent_archive_integration_test.go index ee7d11666..b9da79939 100644 --- a/internal/storage/executionstore/agent_archive_integration_test.go +++ b/internal/storage/executionstore/agent_archive_integration_test.go @@ -161,7 +161,7 @@ tools: } if archivedInteraction.State != executionstore.AgentInteractionStateCanceled || archivedInteraction.ResolvedAt.IsZero() || - archivedInteraction.ResolvedAt.Before(archivedAt) { + archivedInteraction.ResolvedAt.After(archivedAt) { t.Fatalf("archived interaction = %+v, want canceled", archivedInteraction) } toolCall, err := store.Execution().GetToolCall( @@ -174,7 +174,7 @@ tools: if toolCall.State != executionstore.ToolCallStateCompleted || toolCall.Outcome != executionstore.ToolResultOutcomeCanceled || toolCall.CompletedAt == nil || - toolCall.CompletedAt.Before(archivedAt) { + toolCall.CompletedAt.After(archivedAt) { t.Fatalf("archived tool call = %+v, want canceled completion", toolCall) } result, found, err := store.Execution().GetToolCallResultAuthorityByToolCall( @@ -203,3 +203,20 @@ tools: t.Fatalf("archived tool result events = %+v, want one", events) } } + +func TestArchiveAgentStopsAdmittedTurnBeforeFirstAttempt(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "archive_admitted") + ownedInput(t, ctx, f, executionstore.DeliveryModeQueued, "initial") + _, continued, err := f.Store.Execution().AdvanceOwnedAgentWork(ctx, ownedAdvanceInput(f, true)) + require.NoError(t, err) + require.True(t, continued) + _, _, err = f.Store.Execution().ArchiveAgent(ctx, testProjectID, f.AgentID, userPrincipal(f.UserID)) + require.NoError(t, err) + var stops int + require.NoError(t, + f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM agent_events e JOIN agent_inputs i ON i.id=e.agent_input_id WHERE e.agent_id=$1 AND i.control_type='cancel_current'`, + f.AgentID).Scan(&stops)) + require.Equal(t, 1, stops) +} diff --git a/internal/storage/executionstore/agent_config_change_gates_integration_test.go b/internal/storage/executionstore/agent_config_change_gates_integration_test.go index 7d69c3c44..9b77fe8de 100644 --- a/internal/storage/executionstore/agent_config_change_gates_integration_test.go +++ b/internal/storage/executionstore/agent_config_change_gates_integration_test.go @@ -12,7 +12,6 @@ import ( "github.com/omnara-ai/omnara/internal/agentconfig" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" - "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -49,7 +48,7 @@ func (f integrationActivationFixture) launchWithSelectedIntegration( require.NoError(t, integrationstore.LockIntegrationsTx(f.ctx, tx, testProjectID, nil, f.integration.ID)) _, err = txUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: launch.Agent.ID, - }, agentexecution.ExternalAuthority{}) + }, nil) require.NoError(t, err) require.NoError(t, f.store.Integrations().AssignAgentIntegrationConversationTx( f.ctx, tx, testProjectID, launch.Agent.ID, f.integration.ID, diff --git a/internal/storage/executionstore/agent_config_changes_store.go b/internal/storage/executionstore/agent_config_changes_store.go index 3850a1c85..0817ef985 100644 --- a/internal/storage/executionstore/agent_config_changes_store.go +++ b/internal/storage/executionstore/agent_config_changes_store.go @@ -13,7 +13,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -55,9 +54,11 @@ func (s *Store) ChangeAgentConfig( if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil { return ChangeAgentConfigResult{}, errors.New("project and agent are required") } - return storeutil.RetryTransaction(ctx, "change_agent_config", func() (ChangeAgentConfigResult, error) { - return s.changeAgentConfigOnce(ctx, input) - }) + return agentexecution.Retry(ctx, + "change_agent_config", + func(ctx context.Context) (ChangeAgentConfigResult, error) { + return s.changeAgentConfigOnce(ctx, input) + }) } func (s *Store) changeAgentConfigOnce( @@ -317,7 +318,7 @@ func lockAgentForConfigActivationTx( if _, err := unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - agentexecution.IngressAuthority{}); err != nil { + nil); err != nil { return fmt.Errorf("lock agent for config change: %w", err) } return nil diff --git a/internal/storage/executionstore/agent_content_inputs_store.go b/internal/storage/executionstore/agent_content_inputs_store.go index 3849045a3..2ac6e862b 100644 --- a/internal/storage/executionstore/agent_content_inputs_store.go +++ b/internal/storage/executionstore/agent_content_inputs_store.go @@ -57,31 +57,41 @@ func (s *Store) CreateAgentContentInput( if err != nil { return AgentInputRecord{}, nil, false, err } - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentInputRecord{}, nil, false, fmt.Errorf( - "begin create agent content input: %w", - err, - ) + type transactionResult struct { + input AgentInputRecord + contentBlocks json.RawMessage + created bool } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + retried, retryErr := agentexecution.Retry(ctx, + "CreateAgentContentInput", + func(ctx context.Context) (transactionResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return transactionResult{}, fmt.Errorf( + "begin create agent content input: %w", + err, + ) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() - qtx := dbsqlc.New(tx) - result, err := createAgentContentInputTx( - ctx, - unit, - qtx, - input, - contentBlocks, - ) - if err != nil { - return AgentInputRecord{}, nil, false, err - } - if err := unit.Commit(ctx, "create agent content input"); err != nil { - return AgentInputRecord{}, nil, false, err - } - return result.agentInput, result.contentBlocks, result.created, nil + qtx := dbsqlc.New(tx) + result, err := createAgentContentInputTx( + ctx, + unit, + qtx, + input, + contentBlocks, + ) + if err != nil { + return transactionResult{}, err + } + if err := unit.Commit(ctx, "create agent content input"); err != nil { + return transactionResult{}, err + } + return transactionResult{result.agentInput, result.contentBlocks, result.created}, nil + }) + return retried.input, retried.contentBlocks, retried.created, retryErr } func prepareCreateAgentContentInput( @@ -129,7 +139,7 @@ func createAgentContentInputTx( agent, err := unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - agentexecution.IngressAuthority{}, + nil, ) if err != nil { return createAgentContentInputTxResult{}, fmt.Errorf("lock agent for content input: %w", err) diff --git a/internal/storage/executionstore/agent_fixture_integration.go b/internal/storage/executionstore/agent_fixture_integration.go index e3efd1c61..f009c3f59 100644 --- a/internal/storage/executionstore/agent_fixture_integration.go +++ b/internal/storage/executionstore/agent_fixture_integration.go @@ -10,6 +10,7 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/identitystore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" ) @@ -24,7 +25,11 @@ func (s *Store) IntegrationLaunchAgentOnce( ctx context.Context, input LaunchAgentInput, ) (LaunchAgentResult, error) { - return s.launchAgentOnce(ctx, input) + return agentexecution.Retry(ctx, + "IntegrationLaunchAgentOnce", + func(ctx context.Context) (LaunchAgentResult, error) { + return s.launchAgentOnce(ctx, input) + }) } func (s *Store) IntegrationExecuteToolCallOnce( @@ -32,18 +37,26 @@ func (s *Store) IntegrationExecuteToolCallOnce( input ExecuteToolCallInput, command ToolCallCommand, ) (ExecuteToolCallResult, error) { - return s.executeToolCallOnce( - ctx, - input, - func(*ToolCallReader) (ToolCallCommand, error) { return command, nil }, - ) + return agentexecution.Retry(ctx, + "IntegrationExecuteToolCallOnce", + func(ctx context.Context) (ExecuteToolCallResult, error) { + return s.executeToolCallOnce( + ctx, + input, + func(*ToolCallReader) (ToolCallCommand, error) { return command, nil }, + ) + }) } func (s *Store) IntegrationChangeAgentConfigOnce( ctx context.Context, input ChangeAgentConfigInput, ) (ChangeAgentConfigResult, error) { - return s.changeAgentConfigOnce(ctx, input) + return agentexecution.Retry(ctx, + "IntegrationChangeAgentConfigOnce", + func(ctx context.Context) (ChangeAgentConfigResult, error) { + return s.changeAgentConfigOnce(ctx, input) + }) } func (s *Store) IntegrationArchiveAgentOnce( @@ -51,36 +64,62 @@ func (s *Store) IntegrationArchiveAgentOnce( orgID, projectID, agentID uuid.UUID, actor *ActorParams, ) (AgentRecord, []MachineRecord, error) { - _ = orgID - return s.archiveAgentOnce(ctx, projectID, agentID, actor) + type transactionResult struct { + agent AgentRecord + machines []MachineRecord + } + retried, retryErr := agentexecution.Retry(ctx, + "IntegrationArchiveAgentOnce", + func(ctx context.Context) (transactionResult, error) { + _ = orgID + agent, machines, returnErr := s.archiveAgentOnce(ctx, projectID, agentID, actor) + return transactionResult{agent, machines}, returnErr + }) + return retried.agent, retried.machines, retryErr } func (s *Store) IntegrationDeleteMachineOnce( ctx context.Context, input DeleteMachineInput, ) (MachineRecord, error) { - return s.deleteMachineOnce(ctx, input) + return agentexecution.Retry(ctx, + "IntegrationDeleteMachineOnce", + func(ctx context.Context) (MachineRecord, error) { + return s.deleteMachineOnce(ctx, input) + }) } func (s *Store) IntegrationDeleteMachinePoolOnce( ctx context.Context, orgID, poolID uuid.UUID, ) ([]MachineRecord, error) { - return s.deleteMachinePoolOnce(ctx, orgID, poolID) + return agentexecution.Retry(ctx, + "IntegrationDeleteMachinePoolOnce", + func(ctx context.Context) ([]MachineRecord, error) { + return s.deleteMachinePoolOnce(ctx, orgID, poolID) + }) } func (s *Store) IntegrationDeleteProjectMachineGrantOnce( ctx context.Context, orgID, projectID, grantID uuid.UUID, ) (ProjectMachineGrantRecord, error) { - return s.deleteProjectMachineGrantOnce(ctx, orgID, projectID, grantID) + return agentexecution.Retry(ctx, + "IntegrationDeleteProjectMachineGrantOnce", + func(ctx context.Context) (ProjectMachineGrantRecord, error) { + return s.deleteProjectMachineGrantOnce(ctx, orgID, projectID, grantID) + }) } func (s *Store) IntegrationDeleteProjectMachinePoolGrantOnce( ctx context.Context, orgID, projectID, grantID uuid.UUID, ) (DeleteProjectMachinePoolGrantResult, error) { - return s.deleteProjectMachinePoolGrantOnce(ctx, orgID, projectID, grantID) + return agentexecution.Retry(ctx, + "IntegrationDeleteProjectMachinePoolGrantOnce", + func(ctx context.Context) (DeleteProjectMachinePoolGrantResult, error) { + return s.deleteProjectMachinePoolGrantOnce(ctx, orgID, projectID, grantID) + }) } func (s *Store) CreateAgentFixture(ctx context.Context, input AgentFixtureInput) (AgentRecord, error) { @@ -90,54 +129,58 @@ func (s *Store) CreateAgentFixture(ctx context.Context, input AgentFixtureInput) if input.CurrentConfigID == uuid.Nil { return AgentRecord{}, errors.New("current config id is required") } - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentRecord{}, fmt.Errorf("begin create agent fixture: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + return agentexecution.Retry(ctx, + "CreateAgentFixture", + func(ctx context.Context) (AgentRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return AgentRecord{}, fmt.Errorf("begin create agent fixture: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() - qtx := dbsqlc.New(tx) - project, err := loadProjectTx(ctx, qtx, input.ProjectID) - if err != nil { - return AgentRecord{}, err - } - if err := lifecyclelock.EnterActiveProject(ctx, tx, project.OrgID, input.ProjectID); err != nil { - return AgentRecord{}, err - } - config, err := loadAgentConfigTx(ctx, qtx, input.ProjectID, input.CurrentConfigID) - if err != nil { - return AgentRecord{}, err - } - if err := lockAgentConfigForUseTx(ctx, qtx, config); err != nil { - return AgentRecord{}, err - } - if err := lockResourceCreation(ctx, qtx, resourceAgents, input.ProjectID.String()); err != nil { - return AgentRecord{}, err - } - record, _, err := insertAdmittedAgentTx(ctx, unit, qtx, insertAgentInput{ - OrgID: project.OrgID, - ProjectID: input.ProjectID, - Name: input.Name, - CurrentConfigID: input.CurrentConfigID, - }) - if err != nil { - return AgentRecord{}, err - } - if _, err := activateNewAgentConfigTx(ctx, unit, qtx, ActivateAgentConfigInput{ - ProjectID: input.ProjectID, - AgentID: record.ID, - AgentConfigID: input.CurrentConfigID, - ActorType: identitystore.PrincipalTypeSystem, - Reason: "fixture", - IdempotencyKey: "fixture:" + record.ID.String(), - }); err != nil { - return AgentRecord{}, err - } - if err := unit.Commit(ctx, "create agent fixture"); err != nil { - return AgentRecord{}, err - } - return record, nil + qtx := dbsqlc.New(tx) + project, err := loadProjectTx(ctx, qtx, input.ProjectID) + if err != nil { + return AgentRecord{}, err + } + if err := lifecyclelock.EnterActiveProject(ctx, tx, project.OrgID, input.ProjectID); err != nil { + return AgentRecord{}, err + } + config, err := loadAgentConfigTx(ctx, qtx, input.ProjectID, input.CurrentConfigID) + if err != nil { + return AgentRecord{}, err + } + if err := lockAgentConfigForUseTx(ctx, qtx, config); err != nil { + return AgentRecord{}, err + } + if err := lockResourceCreation(ctx, qtx, resourceAgents, input.ProjectID.String()); err != nil { + return AgentRecord{}, err + } + record, _, err := insertAdmittedAgentTx(ctx, unit, qtx, insertAgentInput{ + OrgID: project.OrgID, + ProjectID: input.ProjectID, + Name: input.Name, + CurrentConfigID: input.CurrentConfigID, + }) + if err != nil { + return AgentRecord{}, err + } + if _, err := activateNewAgentConfigTx(ctx, unit, qtx, ActivateAgentConfigInput{ + ProjectID: input.ProjectID, + AgentID: record.ID, + AgentConfigID: input.CurrentConfigID, + ActorType: identitystore.PrincipalTypeSystem, + Reason: "fixture", + IdempotencyKey: "fixture:" + record.ID.String(), + }); err != nil { + return AgentRecord{}, err + } + if err := unit.Commit(ctx, "create agent fixture"); err != nil { + return AgentRecord{}, err + } + return record, nil + }) } func (s *Store) IntegrationBeginMachineWakeOnce( @@ -145,5 +188,9 @@ func (s *Store) IntegrationBeginMachineWakeOnce( orgID, machineID, machinePoolID uuid.UUID, wakeTimeout time.Duration, ) (MachineWakeDisposition, error) { - return s.beginMachineWakeOnce(ctx, orgID, machineID, machinePoolID, wakeTimeout) + return agentexecution.Retry(ctx, + "IntegrationBeginMachineWakeOnce", + func(ctx context.Context) (MachineWakeDisposition, error) { + return s.beginMachineWakeOnce(ctx, orgID, machineID, machinePoolID, wakeTimeout) + }) } diff --git a/internal/storage/executionstore/agent_input_admission_test_helpers_test.go b/internal/storage/executionstore/agent_input_admission_test_helpers_test.go index 0dc54a8e5..6ffb05e90 100644 --- a/internal/storage/executionstore/agent_input_admission_test_helpers_test.go +++ b/internal/storage/executionstore/agent_input_admission_test_helpers_test.go @@ -55,7 +55,7 @@ func admitNextAgentInputAndOpenTurnForTestErr( err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ProjectID: projectID, ID: agentID}, - agentexecution.RuntimeAuthority{AgentID: agentID, + &agentexecution.RuntimeAuthority{AgentID: agentID, RuntimeLockID: runtimeLockID}); err != nil { return executionstore.AdmittedAgentInputTurn{}, false, err } diff --git a/internal/storage/executionstore/agent_inputs_backlog_store.go b/internal/storage/executionstore/agent_inputs_backlog_store.go index 17b34800e..d3151ec43 100644 --- a/internal/storage/executionstore/agent_inputs_backlog_store.go +++ b/internal/storage/executionstore/agent_inputs_backlog_store.go @@ -47,22 +47,27 @@ func (s *Store) DemoteSteeringInputToQueued( func (s *Store) changeBacklog(ctx context.Context, projectID, agentID uuid.UUID, change agentexecution.BacklogChange, cancelInteractions bool) error { - if projectID == uuid.Nil || agentID == uuid.Nil || change.ID == uuid.Nil { - return errors.New("project, agent and input are required") - } - unit, h, err := beginExecution(ctx, s, projectID, agentID) - if err != nil { - return err - } - defer func() { _ = unit.Rollback(ctx) }() - changed, err := h.ChangeBacklog(ctx, change) - if err != nil { - return err - } - if changed && cancelInteractions { - if _, err = h.SupersedeInteractions(ctx, change.ID); err != nil { - return err - } - } - return unit.Commit(ctx, "change input backlog") + _, err := agentexecution.Retry(ctx, + "changeBacklog", + func(ctx context.Context) (struct{}, error) { + if projectID == uuid.Nil || agentID == uuid.Nil || change.ID == uuid.Nil { + return struct{}{}, errors.New("project, agent and input are required") + } + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + changed, err := h.ChangeBacklog(ctx, change) + if err != nil { + return struct{}{}, err + } + if changed && cancelInteractions { + if _, err = h.SupersedeInteractions(ctx, change.ID); err != nil { + return struct{}{}, err + } + } + return struct{}{}, unit.Commit(ctx, "change input backlog") + }) + return err } diff --git a/internal/storage/executionstore/agent_inputs_store.go b/internal/storage/executionstore/agent_inputs_store.go index 847b1a41b..228b847e7 100644 --- a/internal/storage/executionstore/agent_inputs_store.go +++ b/internal/storage/executionstore/agent_inputs_store.go @@ -99,45 +99,54 @@ func (s *Store) ClaimNextAgentWork( ctx context.Context, input ClaimNextAgentWorkInput, ) (ClaimedAgentWork, bool, error) { - if input.WorkerProcessID == uuid.Nil { - return ClaimedAgentWork{}, false, errors.New("worker process id is required") - } - if err := validateAgentRuntimeLockLeaseDuration(input.LeaseDuration); err != nil { - return ClaimedAgentWork{}, false, err - } - unit, err := s.cell.Begin(ctx) - if err != nil { - return ClaimedAgentWork{}, false, err - } - defer func() { _ = unit.Rollback(ctx) }() - route, err := unit.ClaimAgent(ctx) - if errors.Is(err, pgx.ErrNoRows) { - return ClaimedAgentWork{}, false, storeerr.ErrNoClaimableAgentWakeup - } - if err != nil { - return ClaimedAgentWork{}, false, err - } - h, err := unit.Agent(route) - if err != nil { - return ClaimedAgentWork{}, false, err - } - started := time.Now() //nolint:omnaralint // Monotonic lease budget. - work, err := h.Claim(ctx, input.WorkerProcessID, input.LeaseDuration) - if err != nil { - return ClaimedAgentWork{}, false, err - } - if err := applyAdmissionDestination(ctx, unit, route.ProjectID, route.AgentID, work.Admission); err != nil { - return ClaimedAgentWork{}, false, err - } - claim, err := shapeOwnedWork(ctx, unit, route.ProjectID, route.AgentID, work) - if err != nil { - return ClaimedAgentWork{}, false, err - } - claim.LocalLeaseBudgetStartedAt = started - if err = unit.Commit(ctx, "claim agent work"); err != nil { - return ClaimedAgentWork{}, false, err - } - return claim, true, nil + type executionResult struct { + work ClaimedAgentWork + found bool + } + result, err := agentexecution.Retry(ctx, + "ClaimNextAgentWork", + func(ctx context.Context) (executionResult, error) { + if input.WorkerProcessID == uuid.Nil { + return executionResult{ClaimedAgentWork{}, false}, errors.New("worker process id is required") + } + if err := validateAgentRuntimeLockLeaseDuration(input.LeaseDuration); err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + unit, err := s.cell.Begin(ctx) + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + defer func() { _ = unit.Rollback(ctx) }() + route, err := unit.ClaimAgent(ctx) + if errors.Is(err, pgx.ErrNoRows) { + return executionResult{ClaimedAgentWork{}, false}, storeerr.ErrNoClaimableAgentWakeup + } + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + h, err := unit.Agent(route) + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + started := time.Now() //nolint:omnaralint // Monotonic lease budget. + work, err := h.Claim(ctx, input.WorkerProcessID, input.LeaseDuration) + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + if err := applyAdmissionDestination(ctx, unit, route.ProjectID, route.AgentID, work.Admission); err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + claim, err := shapeOwnedWork(ctx, unit, route.ProjectID, route.AgentID, work) + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + claim.LocalLeaseBudgetStartedAt = started + if err = unit.Commit(ctx, "claim agent work"); err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + return executionResult{claim, true}, nil + }) + return result.work, result.found, err } func (s *Store) ListQueuedBacklogInputs( diff --git a/internal/storage/executionstore/agent_interactions_store.go b/internal/storage/executionstore/agent_interactions_store.go index 255a4a68d..7afcd465e 100644 --- a/internal/storage/executionstore/agent_interactions_store.go +++ b/internal/storage/executionstore/agent_interactions_store.go @@ -77,39 +77,44 @@ func (s *Store) CreatePermissionInteraction( ctx context.Context, input CreatePermissionInteractionInput, ) (AgentInteractionRecord, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return AgentInteractionRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - destination, err := captureInteractionDestinationTx(ctx, unit.DB(), input.ProjectID, input.AgentID) - if err != nil { - return AgentInteractionRecord{}, err - } - result, err := h.OpenInteraction(ctx, agentexecution.OpenInteractionInput{ToolRef: agentexecution.ToolRef{ - ID: input.ToolCallID, - RuntimeLockID: input.RuntimeLockID}, - Permission: &input.Request, - Destination: destination}) - if err != nil { - return AgentInteractionRecord{}, err - } - row, err := dbsqlc.New(unit.DB()).GetAgentInteraction(ctx, dbsqlc.GetAgentInteractionParams{ - ProjectID: input.ProjectID, AgentID: input.AgentID, ID: result.ID}) - if err != nil { - return AgentInteractionRecord{}, err - } - if err = unit.Commit(ctx, "open permission"); err != nil { - return AgentInteractionRecord{}, err - } - return agentInteractionRecordFromSQLC(row), nil + return agentexecution.Retry(ctx, + "CreatePermissionInteraction", + func(ctx context.Context) (AgentInteractionRecord, + error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return AgentInteractionRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + destination, err := captureInteractionDestinationTx(ctx, unit.DB(), input.ProjectID, input.AgentID) + if err != nil { + return AgentInteractionRecord{}, err + } + result, err := h.OpenInteraction(ctx, agentexecution.OpenInteractionInput{ToolRef: agentexecution.ToolRef{ + ID: input.ToolCallID, + RuntimeLockID: input.RuntimeLockID}, + Permission: &input.Request, + Destination: destination}) + if err != nil { + return AgentInteractionRecord{}, err + } + row, err := dbsqlc.New(unit.DB()).GetAgentInteraction(ctx, dbsqlc.GetAgentInteractionParams{ + ProjectID: input.ProjectID, AgentID: input.AgentID, ID: result.ID}) + if err != nil { + return AgentInteractionRecord{}, err + } + if err = unit.Commit(ctx, "open permission"); err != nil { + return AgentInteractionRecord{}, err + } + return agentInteractionRecordFromSQLC(row), nil + }) } func (t *toolCallTransaction) createQuestionInteraction( ctx context.Context, input CreateQuestionInteractionInput, ) (AgentInteractionRecord, error) { - plan, err := t.unit.PlanAgentFamily(ctx, t.input.ProjectID, t.input.AgentID, agentexecution.LifecycleAuthority{}) + plan, err := t.unit.PlanAgentFamily(ctx, t.input.ProjectID, t.input.AgentID, nil) if err == nil { err = t.unit.LockAgents(ctx, plan) } @@ -159,19 +164,23 @@ func (s *Store) ResolveAgentInteraction( if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.ID == uuid.Nil { return AgentInteractionRecord{}, errors.New("project, agent, and interaction are required") } - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentInteractionRecord{}, fmt.Errorf("begin resolve agent interaction: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - record, err := resolveExecutionInteraction(ctx, unit, input) - if err != nil { - return AgentInteractionRecord{}, err - } - if err := unit.Commit(ctx, "resolve agent interaction"); err != nil { - return AgentInteractionRecord{}, err - } - return record, nil + return agentexecution.Retry(ctx, + "ResolveAgentInteraction", + func(ctx context.Context) (AgentInteractionRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return AgentInteractionRecord{}, fmt.Errorf("begin resolve agent interaction: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + record, err := resolveExecutionInteraction(ctx, unit, input) + if err != nil { + return AgentInteractionRecord{}, err + } + if err := unit.Commit(ctx, "resolve agent interaction"); err != nil { + return AgentInteractionRecord{}, err + } + return record, nil + }) } func resolveExecutionInteraction( @@ -182,7 +191,7 @@ func resolveExecutionInteraction( if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, AgentID: input.AgentID}}, - agentexecution.ExternalAuthority{}); err != nil { + nil); err != nil { return AgentInteractionRecord{}, err } q := dbsqlc.New(unit.DB()) diff --git a/internal/storage/executionstore/agent_launch_inbox.go b/internal/storage/executionstore/agent_launch_inbox.go index 7b5540396..ba3de1778 100644 --- a/internal/storage/executionstore/agent_launch_inbox.go +++ b/internal/storage/executionstore/agent_launch_inbox.go @@ -9,8 +9,8 @@ import ( "github.com/omnara-ai/omnara/internal/storage/artifactstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -55,9 +55,11 @@ func (s *Store) AdmitInboxLaunchRecipient( if lease.ProjectID == uuid.Nil || lease.ReceiptID == uuid.Nil || lease.Token == uuid.Nil || recipientKey == "" { return LaunchAgentResult{}, storeerr.InvalidRequest(errors.New("inbox lease and recipient are required")) } - return storeutil.RetryTransaction(ctx, "admit_inbox_launch_recipient", func() (LaunchAgentResult, error) { - return s.admitInboxLaunchRecipientOnce(ctx, lease, recipientKey, artifacts) - }) + return agentexecution.Retry(ctx, + "admit_inbox_launch_recipient", + func(ctx context.Context) (LaunchAgentResult, error) { + return s.admitInboxLaunchRecipientOnce(ctx, lease, recipientKey, artifacts) + }) } func (s *Store) admitInboxLaunchRecipientOnce( diff --git a/internal/storage/executionstore/agent_launch_store.go b/internal/storage/executionstore/agent_launch_store.go index 47356ba07..a440f91ba 100644 --- a/internal/storage/executionstore/agent_launch_store.go +++ b/internal/storage/executionstore/agent_launch_store.go @@ -17,7 +17,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -65,9 +64,11 @@ func (s *Store) LaunchAgent( if err != nil { return LaunchAgentResult{}, err } - return storeutil.RetryTransaction(ctx, "launch_agent", func() (LaunchAgentResult, error) { - return s.launchAgentOnce(ctx, input) - }) + return agentexecution.Retry(ctx, + "launch_agent", + func(ctx context.Context) (LaunchAgentResult, error) { + return s.launchAgentOnce(ctx, input) + }) } func validateLaunchAgentInput(input LaunchAgentInput) (LaunchAgentInput, error) { diff --git a/internal/storage/executionstore/agent_runtime_cancellation_store.go b/internal/storage/executionstore/agent_runtime_cancellation_store.go index e5c5bc059..c6279c466 100644 --- a/internal/storage/executionstore/agent_runtime_cancellation_store.go +++ b/internal/storage/executionstore/agent_runtime_cancellation_store.go @@ -25,23 +25,27 @@ type CancelAgentInput struct { const cancelReasonAgentCanceled = "agent_canceled" func (s *Store) CancelAgent(ctx context.Context, input CancelAgentInput) (CancelAgentResult, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return CancelAgentResult{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - actor, err := resolveActorTx(ctx, dbsqlc.New(unit.DB()), input.ProjectID, input.Actor) - if err != nil { - return CancelAgentResult{}, err - } - result, err := h.Cancel(ctx, agentexecution.CancelInput{ActorID: actor, Reason: cancelReasonAgentCanceled, - Message: "The model call was canceled by an explicit agent cancellation."}) - if err != nil { - return CancelAgentResult{}, err - } - if err = unit.Commit(ctx, "cancel agent"); err != nil { - return CancelAgentResult{}, err - } - return CancelAgentResult{Event: executionEvent(input.AgentID, result.Event, events.KindAgentInput), - Affected: result.Changed, RuntimeCancelRequested: result.RuntimeID != uuid.Nil, ActorID: actor}, nil + return agentexecution.Retry(ctx, + "CancelAgent", + func(ctx context.Context) (CancelAgentResult, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return CancelAgentResult{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + actor, err := resolveActorTx(ctx, dbsqlc.New(unit.DB()), input.ProjectID, input.Actor) + if err != nil { + return CancelAgentResult{}, err + } + result, err := h.Cancel(ctx, agentexecution.CancelInput{ActorID: actor, Reason: cancelReasonAgentCanceled, + Message: "The model call was canceled by an explicit agent cancellation."}) + if err != nil { + return CancelAgentResult{}, err + } + if err = unit.Commit(ctx, "cancel agent"); err != nil { + return CancelAgentResult{}, err + } + return CancelAgentResult{Event: executionEvent(input.AgentID, result.Event, events.KindAgentInput), + Affected: result.Changed, RuntimeCancelRequested: result.RuntimeID != uuid.Nil, ActorID: actor}, nil + }) } diff --git a/internal/storage/executionstore/agent_runtime_lock_reaping_store.go b/internal/storage/executionstore/agent_runtime_lock_reaping_store.go index 00e646cb8..a1320fafa 100644 --- a/internal/storage/executionstore/agent_runtime_lock_reaping_store.go +++ b/internal/storage/executionstore/agent_runtime_lock_reaping_store.go @@ -76,20 +76,24 @@ func (s *Store) reapExpiredAgentRuntimeLock( ctx context.Context, projectID, agentID, runtimeLockID uuid.UUID, ) (bool, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return false, fmt.Errorf("begin reap expired agent runtime lock: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() + return agentexecution.Retry(ctx, + "reapExpiredAgentRuntimeLock", + func(ctx context.Context) (bool, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return false, fmt.Errorf("begin reap expired agent runtime lock: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() - reaped, err := reapExpiredAgentRuntimeLockTx(ctx, unit, projectID, agentID, runtimeLockID) - if err != nil || !reaped { - return false, err - } - if err := unit.Commit(ctx, "reap expired agent runtime lock"); err != nil { - return false, err - } - return true, nil + reaped, err := reapExpiredAgentRuntimeLockTx(ctx, unit, projectID, agentID, runtimeLockID) + if err != nil || !reaped { + return false, err + } + if err := unit.Commit(ctx, "reap expired agent runtime lock"); err != nil { + return false, err + } + return true, nil + }) } func reapExpiredAgentRuntimeLockTx( @@ -97,7 +101,7 @@ func reapExpiredAgentRuntimeLockTx( unit *agentexecution.Unit, projectID, agentID, runtimeLockID uuid.UUID, ) (bool, error) { - plan, err := unit.PlanAgentFamily(ctx, projectID, agentID, agentexecution.LifecycleAuthority{}) + plan, err := unit.PlanAgentFamily(ctx, projectID, agentID, nil) var locked bool if err == nil { locked, err = unit.TryLockAgents(ctx, plan) diff --git a/internal/storage/executionstore/agent_runtime_locks_store.go b/internal/storage/executionstore/agent_runtime_locks_store.go index 842cbdc0a..00e6be3fd 100644 --- a/internal/storage/executionstore/agent_runtime_locks_store.go +++ b/internal/storage/executionstore/agent_runtime_locks_store.go @@ -48,21 +48,26 @@ func validateAgentRuntimeLockLeaseDuration(leaseDuration time.Duration) error { func (s *Store) EnsureRuntimeLockActive( ctx context.Context, projectID, agentID, runtimeID uuid.UUID, ) error { - if projectID == uuid.Nil || agentID == uuid.Nil || runtimeID == uuid.Nil { - return errors.New("project, agent, and runtime lock ids are required") - } - unit, h, err := beginExecution(ctx, s, projectID, agentID) - if errors.Is(err, storeerr.ErrNotFound) { - return storeerr.ErrRuntimeLockInactive - } - if err != nil { - return err - } - defer func() { _ = unit.Rollback(ctx) }() - if err := h.FenceRuntime(ctx, runtimeID); err != nil { - return err - } - return unit.Commit(ctx, "runtime lock check") + _, err := agentexecution.Retry(ctx, + "EnsureRuntimeLockActive", + func(ctx context.Context) (struct{}, error) { + if projectID == uuid.Nil || agentID == uuid.Nil || runtimeID == uuid.Nil { + return struct{}{}, errors.New("project, agent, and runtime lock ids are required") + } + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if errors.Is(err, storeerr.ErrNotFound) { + return struct{}{}, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if err := h.FenceRuntime(ctx, runtimeID); err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "runtime lock check") + }) + return err } const RuntimeToolInterruptedMessage = "Tool call was interrupted; its external outcome is unknown." @@ -72,51 +77,60 @@ func (s *Store) RenewAgentRuntimeLock( projectID, agentID, runtimeLockID uuid.UUID, leaseDuration time.Duration, ) (AgentRuntimeLockRenewal, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentRuntimeLockRenewal{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - renewal, err := unit.RenewRuntime( - ctx, - agentexecution.AgentRoute{CellID: unit.CellID(), ProjectID: projectID, AgentID: agentID}, - runtimeLockID, - leaseDuration, - ) - if err != nil { - return AgentRuntimeLockRenewal{}, err - } - row, err := dbsqlc.New(unit.DB()). - GetAgentRuntimeLockForRelease(ctx, - dbsqlc.GetAgentRuntimeLockForReleaseParams{ProjectID: projectID, - AgentID: agentID, - ID: runtimeLockID}) - if err != nil { - return AgentRuntimeLockRenewal{}, err - } - if err = unit.Commit(ctx, "renew runtime"); err != nil { - return AgentRuntimeLockRenewal{}, err - } - return AgentRuntimeLockRenewal{ - RuntimeLock: agentRuntimeLockRecordFromSQLC(row), - LocalLeaseBudgetStartedAt: renewal.LocalStartedAt, - }, nil + return agentexecution.Retry(ctx, + "RenewAgentRuntimeLock", + func(ctx context.Context) (AgentRuntimeLockRenewal, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return AgentRuntimeLockRenewal{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + renewal, err := unit.RenewRuntime( + ctx, + agentexecution.AgentRoute{ProjectID: projectID, AgentID: agentID}, + runtimeLockID, + leaseDuration, + ) + if err != nil { + return AgentRuntimeLockRenewal{}, err + } + row, err := dbsqlc.New(unit.DB()). + GetAgentRuntimeLockForRelease(ctx, + dbsqlc.GetAgentRuntimeLockForReleaseParams{ProjectID: projectID, + AgentID: agentID, + ID: runtimeLockID}) + if err != nil { + return AgentRuntimeLockRenewal{}, err + } + if err = unit.Commit(ctx, "renew runtime"); err != nil { + return AgentRuntimeLockRenewal{}, err + } + return AgentRuntimeLockRenewal{ + RuntimeLock: agentRuntimeLockRecordFromSQLC(row), + LocalLeaseBudgetStartedAt: renewal.LocalStartedAt, + }, nil + }) } func (s *Store) ReleaseAgentRuntimeLock( ctx context.Context, projectID, agentID, runtimeLockID uuid.UUID, ) error { - unit, h, err := beginExecution(ctx, s, projectID, agentID) - if errors.Is(err, storeerr.ErrNotFound) { - return storeerr.ErrRuntimeLockInactive - } - if err != nil { - return err - } - defer func() { _ = unit.Rollback(ctx) }() - if err = h.Release(ctx, runtimeLockID); err != nil { - return err - } - return unit.Commit(ctx, "release runtime") + _, err := agentexecution.Retry(ctx, + "ReleaseAgentRuntimeLock", + func(ctx context.Context) (struct{}, error) { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if errors.Is(err, storeerr.ErrNotFound) { + return struct{}{}, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if err = h.Release(ctx, runtimeLockID); err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "release runtime") + }) + return err } diff --git a/internal/storage/executionstore/agent_runtime_store.go b/internal/storage/executionstore/agent_runtime_store.go index 7c422fe04..0e0f0421b 100644 --- a/internal/storage/executionstore/agent_runtime_store.go +++ b/internal/storage/executionstore/agent_runtime_store.go @@ -409,10 +409,12 @@ func (s *Store) ArchiveAgent( agent AgentRecord machines []MachineRecord } - result, err := storeutil.RetryTransaction(ctx, "archive_agent", func() (archiveAgentResult, error) { - agent, machines, archiveErr := s.archiveAgentOnce(ctx, projectID, agentID, actor) - return archiveAgentResult{agent: agent, machines: machines}, archiveErr - }) + result, err := agentexecution.Retry(ctx, + "archive_agent", + func(ctx context.Context) (archiveAgentResult, error) { + agent, machines, archiveErr := s.archiveAgentOnce(ctx, projectID, agentID, actor) + return archiveAgentResult{agent: agent, machines: machines}, archiveErr + }) return result.agent, result.machines, err } @@ -569,7 +571,7 @@ func lockAgentTreeForArchiveTx( lifecyclelock.AgentRef{ProjectID: root.ProjectID, AgentID: root.ParentAgentID}, ) } - if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, nil); err != nil { return err } lockedTree, err := listActiveAgentTreeTx(ctx, qtx, root.ProjectID, root.ID) diff --git a/internal/storage/executionstore/agent_work_store.go b/internal/storage/executionstore/agent_work_store.go index 158b761b7..066fa1673 100644 --- a/internal/storage/executionstore/agent_work_store.go +++ b/internal/storage/executionstore/agent_work_store.go @@ -5,6 +5,7 @@ import ( "errors" "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -22,30 +23,39 @@ func (s *Store) AdvanceOwnedAgentWork( ctx context.Context, input AdvanceOwnedAgentWorkInput, ) (ClaimedAgentWork, bool, error) { - if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil { - return ClaimedAgentWork{}, false, errors.New("project, agent, and runtime lock ids are required") + type executionResult struct { + work ClaimedAgentWork + found bool } - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if errors.Is(err, storeerr.ErrNotFound) { - return ClaimedAgentWork{}, false, storeerr.ErrRuntimeLockInactive - } - if err != nil { - return ClaimedAgentWork{}, false, err - } - defer func() { _ = unit.Rollback(ctx) }() - work, err := h.Advance(ctx, input.RuntimeLockID, input.AllowModelWork, input.PrepareModel) - if err != nil { - return ClaimedAgentWork{}, false, err - } - if err := applyAdmissionDestination(ctx, unit, input.ProjectID, input.AgentID, work.Admission); err != nil { - return ClaimedAgentWork{}, false, err - } - claim, err := shapeOwnedWork(ctx, unit, input.ProjectID, input.AgentID, work) - if err != nil { - return ClaimedAgentWork{}, false, err - } - if err := unit.Commit(ctx, "advance owned agent work"); err != nil { - return ClaimedAgentWork{}, false, err - } - return claim, !work.Released, nil + result, err := agentexecution.Retry(ctx, + "AdvanceOwnedAgentWork", + func(ctx context.Context) (executionResult, error) { + if input.ProjectID == uuid.Nil || input.AgentID == uuid.Nil || input.RuntimeLockID == uuid.Nil { + return executionResult{ClaimedAgentWork{}, false}, errors.New("project, agent, and runtime lock ids are required") + } + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if errors.Is(err, storeerr.ErrNotFound) { + return executionResult{ClaimedAgentWork{}, false}, storeerr.ErrRuntimeLockInactive + } + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + defer func() { _ = unit.Rollback(ctx) }() + work, err := h.Advance(ctx, input.RuntimeLockID, input.AllowModelWork, input.PrepareModel) + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + if err := applyAdmissionDestination(ctx, unit, input.ProjectID, input.AgentID, work.Admission); err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + claim, err := shapeOwnedWork(ctx, unit, input.ProjectID, input.AgentID, work) + if err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + if err := unit.Commit(ctx, "advance owned agent work"); err != nil { + return executionResult{ClaimedAgentWork{}, false}, err + } + return executionResult{claim, !work.Released}, nil + }) + return result.work, result.found, err } diff --git a/internal/storage/executionstore/context_checkpoints_store.go b/internal/storage/executionstore/context_checkpoints_store.go index 80851eccb..ed7b7ce58 100644 --- a/internal/storage/executionstore/context_checkpoints_store.go +++ b/internal/storage/executionstore/context_checkpoints_store.go @@ -30,37 +30,42 @@ func (s *Store) PublishContextCheckpoint( ctx context.Context, input PublishContextCheckpointInput, ) (ContextCheckpointRecord, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return ContextCheckpointRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - result, err := h.PublishCheckpoint( - ctx, - agentexecution.PublishCheckpointInput{RuntimeLockID: input.RuntimeLockID, - ContextID: input.ModelCallContextID, Summary: input.Summary, Evidence: agentexecution.ModelEvidence{ - APIFormat: input.APIFormat, APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, - ResponseID: input.ProviderResponseID, - Usage: input.Usage, - Cost: input.ProviderReportedCostUSD, - Metadata: input.ProviderMetadata}}, - ) - if err != nil { - return ContextCheckpointRecord{}, err - } - if !result.Created { - return ContextCheckpointRecord{}, storeerr.ErrStateTransitionConflict - } - row, err := dbsqlc.New(unit.DB()). - GetContextCheckpoint(ctx, - dbsqlc.GetContextCheckpointParams{ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: result.ID}) - if err != nil { - return ContextCheckpointRecord{}, err - } - if err = unit.Commit(ctx, "publish context checkpoint"); err != nil { - return ContextCheckpointRecord{}, err - } - return contextCheckpointRecordFromGetSQLC(row), nil + return agentexecution.Retry(ctx, + "PublishContextCheckpoint", + func(ctx context.Context) (ContextCheckpointRecord, + error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return ContextCheckpointRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.PublishCheckpoint( + ctx, + agentexecution.PublishCheckpointInput{RuntimeLockID: input.RuntimeLockID, + ContextID: input.ModelCallContextID, Summary: input.Summary, Evidence: agentexecution.ModelEvidence{ + APIFormat: input.APIFormat, APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, + ResponseID: input.ProviderResponseID, + Usage: input.Usage, + Cost: input.ProviderReportedCostUSD, + Metadata: input.ProviderMetadata}}, + ) + if err != nil { + return ContextCheckpointRecord{}, err + } + if !result.Created { + return ContextCheckpointRecord{}, storeerr.ErrStateTransitionConflict + } + row, err := dbsqlc.New(unit.DB()). + GetContextCheckpoint(ctx, + dbsqlc.GetContextCheckpointParams{ProjectID: input.ProjectID, + AgentID: input.AgentID, + ID: result.ID}) + if err != nil { + return ContextCheckpointRecord{}, err + } + if err = unit.Commit(ctx, "publish context checkpoint"); err != nil { + return ContextCheckpointRecord{}, err + } + return contextCheckpointRecordFromGetSQLC(row), nil + }) } diff --git a/internal/storage/executionstore/context_checkpoints_terminal_failure.go b/internal/storage/executionstore/context_checkpoints_terminal_failure.go index b8c7d9c28..2eb0ed5d2 100644 --- a/internal/storage/executionstore/context_checkpoints_terminal_failure.go +++ b/internal/storage/executionstore/context_checkpoints_terminal_failure.go @@ -33,25 +33,30 @@ func (s *Store) RecordTerminalCompactionFailure( ctx context.Context, input RecordTerminalCompactionFailureInput, ) error { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return err - } - defer func() { _ = unit.Rollback(ctx) }() - _, err = h.FailModel( - ctx, - agentexecution.ModelFailure{ContextID: input.ModelCallContextID, RuntimeLockID: input.RuntimeLockID, - ErrorKind: input.ErrorKind, - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: input.ErrorDetails, - ServedModel: input.ServedProviderModelSlug, - Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, - APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, - Usage: input.Usage, Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}}, - ) - if err != nil { - return err - } - return unit.Commit(ctx, "record terminal compaction failure") + _, err := agentexecution.Retry(ctx, + "RecordTerminalCompactionFailure", + func(ctx context.Context) (struct{}, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + _, err = h.FailModel( + ctx, + agentexecution.ModelFailure{ContextID: input.ModelCallContextID, RuntimeLockID: input.RuntimeLockID, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + ServedModel: input.ServedProviderModelSlug, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, + APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, + Usage: input.Usage, Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}}, + ) + if err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "record terminal compaction failure") + }) + return err } diff --git a/internal/storage/executionstore/cron_trigger_agent_inputs_store.go b/internal/storage/executionstore/cron_trigger_agent_inputs_store.go index 9bed4ffe9..0edded525 100644 --- a/internal/storage/executionstore/cron_trigger_agent_inputs_store.go +++ b/internal/storage/executionstore/cron_trigger_agent_inputs_store.go @@ -7,6 +7,7 @@ import ( "fmt" "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -51,26 +52,30 @@ func (s *Store) CreateCronTriggerAgentInput(ctx context.Context, input CreateCro return err } - unit, err := s.cell.Begin(ctx) - if err != nil { - return fmt.Errorf("begin cron trigger input delivery: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + _, retryErr := agentexecution.Retry(ctx, + "CreateCronTriggerAgentInput", + func(ctx context.Context) (struct{}, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return struct{}{}, fmt.Errorf("begin cron trigger input delivery: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + qtx := dbsqlc.New(tx) + if _, err := createAgentContentInputTx(ctx, unit, qtx, contentInput, contentBlocks); err != nil { + return struct{}{}, err + } + if err := completeCronTriggerFiringTx(ctx, qtx, CompleteCronTriggerFiringInput{ + ProjectID: trigger.ProjectID, + TriggerID: trigger.TriggerID, + ClaimToken: trigger.ClaimToken, + Fired: true, + }); err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "deliver cron trigger input") + }) + return retryErr - qtx := dbsqlc.New(tx) - // Input creation locks the agent before completion locks the trigger, matching - // the lock order used when archiving an agent and deleting its cron triggers. - if _, err := createAgentContentInputTx(ctx, unit, qtx, contentInput, contentBlocks); err != nil { - return err - } - if err := completeCronTriggerFiringTx(ctx, qtx, CompleteCronTriggerFiringInput{ - ProjectID: trigger.ProjectID, - TriggerID: trigger.TriggerID, - ClaimToken: trigger.ClaimToken, - Fired: true, - }); err != nil { - return err - } - return unit.Commit(ctx, "deliver cron trigger input") } diff --git a/internal/storage/executionstore/execution_mutation.go b/internal/storage/executionstore/execution_mutation.go index 7cb2d58f7..b1a7e41e4 100644 --- a/internal/storage/executionstore/execution_mutation.go +++ b/internal/storage/executionstore/execution_mutation.go @@ -5,6 +5,7 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" ) func beginExecution( @@ -16,7 +17,7 @@ func beginExecution( if err != nil { return nil, nil, err } - plan, err := unit.PlanAgentFamily(ctx, projectID, agentID, agentexecution.LifecycleAuthority{}) + plan, err := unit.PlanAgents(ctx, []lifecyclelock.AgentRef{{ProjectID: projectID, AgentID: agentID}}, nil) if err == nil { err = unit.LockAgents(ctx, plan) } diff --git a/internal/storage/executionstore/execution_unit_lock_order_integration_test.go b/internal/storage/executionstore/execution_unit_lock_order_integration_test.go index 91d900271..e51f94595 100644 --- a/internal/storage/executionstore/execution_unit_lock_order_integration_test.go +++ b/internal/storage/executionstore/execution_unit_lock_order_integration_test.go @@ -78,31 +78,19 @@ func assertExecutionAgentOrder( ) } -func TestModelMutationPlansParentBeforeRuntime(t *testing.T) { +func TestModelMutationsFenceWithoutParentLock(t *testing.T) { for _, parentFirst := range []bool{true, false} { t.Run(fmt.Sprintf("parent_first_%t", parentFirst), func(t *testing.T) { f, parent, child := executionLockOrderAgents(t, parentFirst) - first, second := parent, child - if !parentFirst { - first, second = child, parent - } runtime, source := uuid.New(), uuid.New() for _, operation := range []string{ - "claim normal", "claim compaction", "retry", "prepare", + "claim compaction", "retry", "prepare", "begin compaction", "replace compaction", "terminal compaction", } { t.Run(operation, func(t *testing.T) { - assertExecutionAgentOrder(t, f, first, second, func(ctx context.Context) error { + assertExecutionAgentOrder(t, f, child, parent, func(ctx context.Context) error { s := f.Store.Execution() switch operation { - case "claim normal": - _, err := s.PrepareNormalModelCall( - ctx, - executionstore.PrepareNormalModelCallInput{ - ProjectID: testProjectID, AgentID: child, RuntimeLockID: runtime, - OpeningInputIDs: []uuid.UUID{source}}, - ) - return err case "claim compaction": _, err := s.ClaimCompactionModelCall( ctx, @@ -256,16 +244,16 @@ func TestSubagentLaunchTakesQuotaBeforeParentAndMachines(t *testing.T) { require.NoError(t, integrationdb.Await(t, done, "subagent launch after quota").Err) } -func TestPreparationFencesLeaseAfterParentWait(t *testing.T) { +func TestPreparationFencesLeaseAfterAgentWait(t *testing.T) { t.Parallel() ctx := t.Context() - f, parent, child := executionLockOrderAgents(t, true) + f, _, child := executionLockOrderAgents(t, true) runtime, err := f.Store.Execution().AcquireAgentRuntimeLock( ctx, testProjectID, child, testWorkerProcessID, time.Minute, ) require.NoError(t, err) blocker := integrationdb.BeginTx(t, ctx, f.Store.pool) - _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, parent) + _, err = blocker.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, child) require.NoError(t, err) done := integrationdb.RunAsyncError(func() error { _, err := f.Store.Execution().PrepareNormalModelCall(ctx, executionstore.PrepareNormalModelCallInput{ diff --git a/internal/storage/executionstore/execution_unit_webhooks_integration_test.go b/internal/storage/executionstore/execution_unit_webhooks_integration_test.go index 2efdfaba1..b177b7c3f 100644 --- a/internal/storage/executionstore/execution_unit_webhooks_integration_test.go +++ b/internal/storage/executionstore/execution_unit_webhooks_integration_test.go @@ -9,11 +9,10 @@ import ( "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage" - "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/stretchr/testify/require" ) -func TestLifecycleUnitEnqueuesWebhooksAtomically(t *testing.T) { +func TestLifecycleDeletionSuppressesWebhooks(t *testing.T) { for _, scope := range []string{"project", "organization"} { for _, fail := range []bool{false, true} { t.Run(fmt.Sprintf("%s/fail_%t", scope, fail), func(t *testing.T) { @@ -21,7 +20,7 @@ func TestLifecycleUnitEnqueuesWebhooksAtomically(t *testing.T) { f := newProcessDaemonFixture(t, ctx, "lifecycle-webhooks") toolID := createToolCallForProcessTestWithPermission(t, ctx, f, "lifecycle-question", "ask_question", true) enableEventWebhook(t, f, []string{"agent_input", "tool_result", "tool_call_update"}) - question := createQuestionInteractionForTest(t, ctx, f, toolID) + createQuestionInteractionForTest(t, ctx, f, toolID) publisher := &recordingPostCommitPublisher{} f.Store = newIntegrationStore(f.Store.pool, storage.WithPostCommitPublisher(publisher)) _, err := f.Store.pool.Exec(ctx, `DELETE FROM event_webhook_deliveries WHERE agent_id=$1`, f.AgentID) @@ -43,33 +42,15 @@ FOR EACH ROW EXECUTE FUNCTION reject_unit_webhook();`) return err } err = remove(ctx) - if fail { - require.ErrorContains(t, err, "EnqueueEventWebhookDelivery") - } else { - require.NoError(t, err) - } + require.NoError(t, err) var state string require.NoError(t, f.Store.pool.QueryRow(ctx, `SELECT state FROM agents WHERE id=$1`, f.AgentID).Scan(&state)) var count int require.NoError(t, f.Store.pool.QueryRow(ctx, `SELECT count(*) FROM event_webhook_deliveries WHERE agent_id=$1`, f.AgentID).Scan(&count)) - if fail { - require.Equal(t, "active", state) - require.Zero(t, count) - require.Empty(t, publisher.intents) - var interactionState string - require.NoError(t, f.Store.pool.QueryRow(ctx, - `SELECT state FROM agent_interactions WHERE id=$1`, question.ID).Scan(&interactionState)) - require.Equal(t, string(executionstore.AgentInteractionStateOpen), interactionState) - return - } require.Equal(t, "archived", state) - require.Positive(t, count) - var completed int - require.NoError(t, f.Store.pool.QueryRow(ctx, - `SELECT count(*) FROM event_webhook_deliveries WHERE agent_id=$1 AND tool_call_id=$2 AND tool_state='completed'`, f.AgentID, toolID).Scan(&completed)) - require.Equal(t, 1, completed) + require.Zero(t, count) var published bool for _, intent := range publisher.intents { if update, ok := intent.(notifications.ToolCallUpdatedCommitted); ok && diff --git a/internal/storage/executionstore/export_integration_test.go b/internal/storage/executionstore/export_integration_test.go index ff005be51..78e57f26e 100644 --- a/internal/storage/executionstore/export_integration_test.go +++ b/internal/storage/executionstore/export_integration_test.go @@ -31,7 +31,7 @@ type ModelWorkSeed struct { } func (s *Store) IntegrationBeginUnit(ctx context.Context) (*agentexecution.Unit, error) { - return s.cell.Begin(ctx) + return agentexecution.BeginTestUnit(ctx, s.cell) } func (s *Store) AcquireAgentRuntimeLock( @@ -42,53 +42,63 @@ func (s *Store) AcquireAgentRuntimeLock( if err := validateAgentRuntimeLockLeaseDuration(lease); err != nil { return AgentRuntimeLockRecord{}, err } - unit, h, err := beginExecution(ctx, s, projectID, agentID) - if err != nil { - return AgentRuntimeLockRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - snapshot, err := h.LoadExecution(ctx) - if err != nil { - return AgentRuntimeLockRecord{}, err - } - if snapshot.View.State != agentexecution.AgentActive { - return AgentRuntimeLockRecord{}, storeerr.ErrAgentNotAdvanceable - } - var record AgentRuntimeLockRecord - err = unit.DB(). - QueryRow(ctx, - `INSERT INTO agent_runtime_locks(agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + return agentexecution.Retry(ctx, + "AcquireAgentRuntimeLock", + func(ctx context.Context) (AgentRuntimeLockRecord, error) { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if err != nil { + return AgentRuntimeLockRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return AgentRuntimeLockRecord{}, err + } + if snapshot.View.State != agentexecution.AgentActive { + return AgentRuntimeLockRecord{}, storeerr.ErrAgentNotAdvanceable + } + var record AgentRuntimeLockRecord + err = unit.DB(). + QueryRow(ctx, + `INSERT INTO agent_runtime_locks(agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) VALUES($1,$2,statement_timestamp(),statement_timestamp(),statement_timestamp()+$3::bigint*interval '1 microsecond') RETURNING id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at,cancel_requested_at`, - agentID, - workerID, - lease.Microseconds()). - Scan(&record.ID, - &record.AgentID, - &record.WorkerProcessID, - &record.StartedAt, - &record.RenewedAt, - &record.LeaseExpiresAt, - &record.CancelRequestedAt) - if err != nil { - return AgentRuntimeLockRecord{}, err - } - if _, err = h.Repair(ctx); err != nil { - return AgentRuntimeLockRecord{}, err - } - return record, unit.Commit(ctx, "seed runtime") + agentID, + workerID, + lease.Microseconds()). + Scan(&record.ID, + &record.AgentID, + &record.WorkerProcessID, + &record.StartedAt, + &record.RenewedAt, + &record.LeaseExpiresAt, + &record.CancelRequestedAt) + if err != nil { + return AgentRuntimeLockRecord{}, err + } + if _, err = h.Repair(ctx); err != nil { + return AgentRuntimeLockRecord{}, err + } + return record, unit.Commit(ctx, "seed runtime") + }) } func (s *Store) RepairExecutionForTest(ctx context.Context, projectID, agentID uuid.UUID) error { - unit, h, err := beginExecution(ctx, s, projectID, agentID) - if err != nil { - return err - } - defer func() { _ = unit.Rollback(ctx) }() - if _, err = h.Repair(ctx); err != nil { - return err - } - return unit.Commit(ctx, "repair execution") + _, retryErr := agentexecution.Retry(ctx, + "RepairExecutionForTest", + func(ctx context.Context) (struct{}, error) { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if _, err = h.Repair(ctx); err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "repair execution") + }) + return retryErr + } func (s *Store) DeleteAgentWakeup(ctx context.Context, projectID, agentID uuid.UUID) error { @@ -105,27 +115,36 @@ func (s *Store) NextAgentModelWork( ctx context.Context, projectID, agentID uuid.UUID, ) (ModelWorkSeed, bool, error) { - unit, h, err := beginExecution(ctx, s, projectID, agentID) - if err != nil { - return ModelWorkSeed{}, false, err - } - defer func() { _ = unit.Rollback(ctx) }() - snapshot, err := h.LoadExecution(ctx) - if err != nil { - return ModelWorkSeed{}, false, err - } - work := snapshot.Selection.Model - if work == nil { - return ModelWorkSeed{}, false, nil - } - return ModelWorkSeed{ - Kind: ModelWorkKind(work.Kind), - ModelCallContextID: work.SourceContextID, - SourceModelOutputID: work.SourceOutputID, - TurnID: work.TurnID, - InputIDs: work.Opening.InputIDs, - OpeningEventSequence: work.Opening.EventSequence, - }, true, nil + type transactionResult struct { + work ModelWorkSeed + found bool + } + retried, retryErr := agentexecution.Retry(ctx, + "NextAgentModelWork", + func(ctx context.Context) (transactionResult, error) { + unit, h, err := beginExecution(ctx, s, projectID, agentID) + if err != nil { + return transactionResult{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + snapshot, err := h.LoadExecution(ctx) + if err != nil { + return transactionResult{}, err + } + work := snapshot.Selection.Model + if work == nil { + return transactionResult{}, nil + } + return transactionResult{ModelWorkSeed{ + Kind: ModelWorkKind(work.Kind), + ModelCallContextID: work.SourceContextID, + SourceModelOutputID: work.SourceOutputID, + TurnID: work.TurnID, + InputIDs: work.Opening.InputIDs, + OpeningEventSequence: work.Opening.EventSequence, + }, true}, nil + }) + return retried.work, retried.found, retryErr } func (s *Store) GetToolCallResultAuthorityByToolCall( @@ -228,7 +247,7 @@ func IntegrationEnsureRuntimeLockActiveTx( ) error { if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: projectID, AgentID: agentID}}, - agentexecution.RuntimeAuthority{AgentID: agentID, RuntimeLockID: runtimeID}, + &agentexecution.RuntimeAuthority{AgentID: agentID, RuntimeLockID: runtimeID}, ); err != nil { return err } @@ -349,7 +368,7 @@ func IntegrationRenewAgentRuntimeLockTx( ) (AgentRuntimeLockRenewal, error) { renewal, err := unit.RenewRuntime( ctx, - agentexecution.AgentRoute{CellID: unit.CellID(), ProjectID: projectID, AgentID: agentID}, + agentexecution.AgentRoute{ProjectID: projectID, AgentID: agentID}, runtimeID, lease, ) diff --git a/internal/storage/executionstore/inbox_input_store.go b/internal/storage/executionstore/inbox_input_store.go index 1c2ad57be..d4c8ef79a 100644 --- a/internal/storage/executionstore/inbox_input_store.go +++ b/internal/storage/executionstore/inbox_input_store.go @@ -12,7 +12,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -48,9 +47,11 @@ func (s *Store) AdmitInboxInputRecipient( if lease.ProjectID == uuid.Nil || lease.ReceiptID == uuid.Nil || lease.Token == uuid.Nil || recipientKey == "" { return InboxInputResult{}, storeerr.InvalidRequest(errors.New("inbox lease and recipient are required")) } - return storeutil.RetryTransaction(ctx, "admit_inbox_input_recipient", func() (InboxInputResult, error) { - return s.admitInboxInputRecipientOnce(ctx, lease, recipientKey, artifacts) - }) + return agentexecution.Retry(ctx, + "admit_inbox_input_recipient", + func(ctx context.Context) (InboxInputResult, error) { + return s.admitInboxInputRecipientOnce(ctx, lease, recipientKey, artifacts) + }) } func (s *Store) admitInboxInputRecipientOnce( @@ -111,7 +112,7 @@ func (s *Store) admitInboxInputRecipientOnce( } if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ ProjectID: lease.ProjectID, AgentID: recipient.AgentID, - }}, agentexecution.IngressAuthority{}); err != nil { + }}, nil); err != nil { return InboxInputResult{}, err } outcome, err := resolveInboxInputOutcome(ctx, q, recipient) diff --git a/internal/storage/executionstore/integration_input_origin.go b/internal/storage/executionstore/integration_input_origin.go index 5b69c7781..f5305127a 100644 --- a/internal/storage/executionstore/integration_input_origin.go +++ b/internal/storage/executionstore/integration_input_origin.go @@ -98,7 +98,7 @@ func (s *Store) admitOriginContentTx( if err := unit.LockAgentRefs( ctx, []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, AgentID: input.AgentID}}, - agentexecution.IngressAuthority{}, + nil, ); err != nil { return InboxInputResult{}, err } diff --git a/internal/storage/executionstore/integration_interactions_integration_test.go b/internal/storage/executionstore/integration_interactions_integration_test.go index 7fd12d788..82ab3a564 100644 --- a/internal/storage/executionstore/integration_interactions_integration_test.go +++ b/internal/storage/executionstore/integration_interactions_integration_test.go @@ -18,7 +18,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/identitystore" "github.com/omnara-ai/omnara/internal/storage/integrationstore" - "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -235,7 +234,7 @@ func (f integrationInteractionFixture) selectOrigin( defer func() { _ = txUnit.Rollback(f.ctx) }() _, err := txUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, - }, agentexecution.ExternalAuthority{}) + }, nil) require.NoError(t, err) selection, err := executionstore.IntegrationSelectInteractionDestinationForOriginTx( f.ctx, txUnit, testProjectID, f.process.AgentID, targetID, @@ -536,7 +535,7 @@ func TestIntegrationInteractionsRejectForeignCallbackAndTarget(t *testing.T) { defer func() { _ = txUnit.Rollback(f.ctx) }() _, err := txUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, - }, agentexecution.ExternalAuthority{}) + }, nil) require.NoError(t, err) selection, err := executionstore.IntegrationSelectInteractionDestinationForOriginTx( f.ctx, txUnit, testProjectID, f.process.AgentID, foreign.ID, @@ -612,7 +611,7 @@ func TestIntegrationInteractionsCaptureUsesLockedCurrentSelectionWithoutIntegrat defer func() { _ = selectionTxUnit.Rollback(f.ctx) }() _, err = selectionTxUnit.LockAgent(f.ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: testProjectID, ID: f.process.AgentID, - }, agentexecution.ExternalAuthority{}) + }, nil) require.NoError(t, err) done := integrationdb.RunAsync(func() (executionstore.AgentInteractionRecord, error) { return f.store.Execution(). diff --git a/internal/storage/executionstore/integration_interactions_store.go b/internal/storage/executionstore/integration_interactions_store.go index 533df2f06..3dc4bc6e8 100644 --- a/internal/storage/executionstore/integration_interactions_store.go +++ b/internal/storage/executionstore/integration_interactions_store.go @@ -44,7 +44,7 @@ func (s *Store) ReconcileInteractionSelectionInUnit( if err := unit.LockAgentRefs( ctx, []lifecyclelock.AgentRef{{ProjectID: projectID, AgentID: agentID}}, - agentexecution.IngressAuthority{}, + nil, ); err != nil { return InteractionSelection{}, err } @@ -556,63 +556,71 @@ func (s *Store) GetInteractionForHandlerCallback( func (s *Store) ResolveAgentInteractionFromHandler( ctx context.Context, input ResolveAgentInteractionFromHandlerInput, ) (AgentInteractionRecord, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentInteractionRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - _, destination, err := lockAuthorizedInteractionDestination( - ctx, - unit, - input.ProjectID, - input.AgentID, - input.ID, - ) - if err != nil { - return AgentInteractionRecord{}, err - } - integration, err := s.integrations.GetIntegrationByIDTx(ctx, tx, destination.IntegrationID) - if err != nil { - return AgentInteractionRecord{}, err - } - if integration.ProjectID != input.ProjectID || - (input.SourceSetupRevision != 0 && integration.SetupRevision != input.SourceSetupRevision) { - return AgentInteractionRecord{}, storeerr.ErrUnauthorized - } - if destination.IntegrationID != input.IntegrationID || - destination.IntegrationKind != input.IntegrationKind || - destination.Address != input.Address || - (input.IntegrationTargetID != uuid.Nil && input.IntegrationTargetID != destination.IntegrationTargetID) { - return AgentInteractionRecord{}, storeerr.ErrUnauthorized - } - if err := validateIntegrationInputActor(integration, input.Actor); err != nil { - return AgentInteractionRecord{}, err - } - record, err := resolveExecutionInteraction( - ctx, - unit, - input.ResolveAgentInteractionInput, - ) - if err != nil { - return AgentInteractionRecord{}, err - } - if err := unit.Commit(ctx, "resolve interaction from handler"); err != nil { - return AgentInteractionRecord{}, err - } - return record, nil + return agentexecution.Retry(ctx, + "ResolveAgentInteractionFromHandler", + func(ctx context.Context) (AgentInteractionRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return AgentInteractionRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + _, destination, err := lockAuthorizedInteractionDestination( + ctx, + unit, + input.ProjectID, + input.AgentID, + input.ID, + ) + if err != nil { + return AgentInteractionRecord{}, err + } + integration, err := s.integrations.GetIntegrationByIDTx(ctx, tx, destination.IntegrationID) + if err != nil { + return AgentInteractionRecord{}, err + } + if integration.ProjectID != input.ProjectID || + (input.SourceSetupRevision != 0 && integration.SetupRevision != input.SourceSetupRevision) { + return AgentInteractionRecord{}, storeerr.ErrUnauthorized + } + if destination.IntegrationID != input.IntegrationID || + destination.IntegrationKind != input.IntegrationKind || + destination.Address != input.Address || + (input.IntegrationTargetID != uuid.Nil && input.IntegrationTargetID != destination.IntegrationTargetID) { + return AgentInteractionRecord{}, storeerr.ErrUnauthorized + } + if err := validateIntegrationInputActor(integration, input.Actor); err != nil { + return AgentInteractionRecord{}, err + } + record, err := resolveExecutionInteraction( + ctx, + unit, + input.ResolveAgentInteractionInput, + ) + if err != nil { + return AgentInteractionRecord{}, err + } + if err := unit.Commit(ctx, "resolve interaction from handler"); err != nil { + return AgentInteractionRecord{}, err + } + return record, nil + }) } func (s *Store) GetAgentInteractionForPresentation( ctx context.Context, projectID, agentID, id uuid.UUID, ) (AgentInteractionRecord, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentInteractionRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - record, _, err := lockAuthorizedInteractionDestination(ctx, unit, projectID, agentID, id) - return record, err + return agentexecution.Retry(ctx, + "GetAgentInteractionForPresentation", + func(ctx context.Context) (AgentInteractionRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return AgentInteractionRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + record, _, err := lockAuthorizedInteractionDestination(ctx, unit, projectID, agentID, id) + return record, err + }) } func lockAuthorizedInteractionDestination( @@ -654,7 +662,7 @@ func lockAuthorizedInteractionDestination( } if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: projectID, ID: agentID, - }, agentexecution.ExternalAuthority{}); err != nil { + }, nil); err != nil { return AgentInteractionRecord{}, nil, err } _, contract, err := loadInteractionContract(ctx, q, projectID, agentID) @@ -691,62 +699,66 @@ func (s *Store) RecordInteractionPresentationReceipt( if err := validateInteractionObject(input.Receipt, InteractionReceiptMaxBytes); err != nil { return AgentInteractionRecord{}, storeerr.InvalidRequest(err) } - unit, err := s.cell.Begin(ctx) - if err != nil { - return AgentInteractionRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - q := dbsqlc.New(tx) - if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, ID: input.AgentID, - }, agentexecution.ExternalAuthority{}); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return AgentInteractionRecord{}, storeerr.ErrNotFound - } - return AgentInteractionRecord{}, err - } - params := dbsqlc.GetAgentInteractionParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - } - row, err := q.GetAgentInteraction(ctx, params) - if errors.Is(err, pgx.ErrNoRows) { - return AgentInteractionRecord{}, storeerr.ErrNotFound - } - if err != nil { - return AgentInteractionRecord{}, err - } - record := agentInteractionRecordFromSQLC(row) - destination, err := record.CapturedDestination() - if err != nil { - return AgentInteractionRecord{}, err - } - if destination == nil || *destination != input.Destination { - return AgentInteractionRecord{}, storeerr.ErrUnauthorized - } - if len(record.PresentationReceipt) != 0 { - if !sameJSON(record.PresentationReceipt, input.Receipt) { - return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict - } - return record, nil - } - h, err := unit.Handle(input.ProjectID, input.AgentID) - if err != nil { - return AgentInteractionRecord{}, err - } - if _, err = h.RecordPresentation(ctx, input.ID, record.Destination, input.Receipt); err != nil { - return AgentInteractionRecord{}, err - } - row, err = q.GetAgentInteraction(ctx, params) - if err != nil { - return AgentInteractionRecord{}, err - } - if err := unit.Commit(ctx, "RecordInteractionPresentationReceipt"); err != nil { - return AgentInteractionRecord{}, err - } - return agentInteractionRecordFromSQLC(row), nil + return agentexecution.Retry(ctx, + "RecordInteractionPresentationReceipt", + func(ctx context.Context) (AgentInteractionRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return AgentInteractionRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + q := dbsqlc.New(tx) + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: input.ProjectID, ID: input.AgentID, + }, nil); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return AgentInteractionRecord{}, storeerr.ErrNotFound + } + return AgentInteractionRecord{}, err + } + params := dbsqlc.GetAgentInteractionParams{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + ID: input.ID, + } + row, err := q.GetAgentInteraction(ctx, params) + if errors.Is(err, pgx.ErrNoRows) { + return AgentInteractionRecord{}, storeerr.ErrNotFound + } + if err != nil { + return AgentInteractionRecord{}, err + } + record := agentInteractionRecordFromSQLC(row) + destination, err := record.CapturedDestination() + if err != nil { + return AgentInteractionRecord{}, err + } + if destination == nil || *destination != input.Destination { + return AgentInteractionRecord{}, storeerr.ErrUnauthorized + } + if len(record.PresentationReceipt) != 0 { + if !sameJSON(record.PresentationReceipt, input.Receipt) { + return AgentInteractionRecord{}, storeerr.ErrIdempotencyConflict + } + return record, nil + } + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return AgentInteractionRecord{}, err + } + if _, err = h.RecordPresentation(ctx, input.ID, record.Destination, input.Receipt); err != nil { + return AgentInteractionRecord{}, err + } + row, err = q.GetAgentInteraction(ctx, params) + if err != nil { + return AgentInteractionRecord{}, err + } + if err := unit.Commit(ctx, "RecordInteractionPresentationReceipt"); err != nil { + return AgentInteractionRecord{}, err + } + return agentInteractionRecordFromSQLC(row), nil + }) } func interactionSelectionForOrigin( diff --git a/internal/storage/executionstore/interaction_selection_queries_integration_test.go b/internal/storage/executionstore/interaction_selection_queries_integration_test.go index 56f8cad15..a91c2e678 100644 --- a/internal/storage/executionstore/interaction_selection_queries_integration_test.go +++ b/internal/storage/executionstore/interaction_selection_queries_integration_test.go @@ -8,7 +8,6 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/stretchr/testify/require" ) @@ -30,7 +29,7 @@ func TestInteractionSelectionOriginPrecedesBackgroundActor(t *testing.T) { _, err = unit.LockAgent( f.ctx, dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: f.process.AgentID}, - agentexecution.ExternalAuthority{}, + nil, ) require.NoError(t, err) q := dbsqlc.New(unit.DB()) diff --git a/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go b/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go index 40411d991..c8cfe1e77 100644 --- a/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go +++ b/internal/storage/executionstore/kernel_event_authority_helpers_integration_test.go @@ -17,7 +17,6 @@ import ( "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/toolcatalog" "github.com/stretchr/testify/require" @@ -692,7 +691,7 @@ func claimSentCompactionForRangeTest( _, err = unit.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: fixture.AgentID}, - agentexecution.ExternalAuthority{}, + nil, ) require.NoError(t, err) admitted, err := executionstore.IntegrationAdmitInputs(ctx, unit, testProjectID, fixture.AgentID) diff --git a/internal/storage/executionstore/lifecycle_residual_integration_test.go b/internal/storage/executionstore/lifecycle_residual_integration_test.go new file mode 100644 index 000000000..2741e0e8d --- /dev/null +++ b/internal/storage/executionstore/lifecycle_residual_integration_test.go @@ -0,0 +1,96 @@ +//go:build integration + +package executionstore_test + +import ( + "testing" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/processcmd" + "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/stretchr/testify/require" +) + +func TestOrganizationDeletionIncludesDeletedProjectPoolAgents(t *testing.T) { + ctx := t.Context() + f := newMachineLifecycleLockOrderFixture(t, ctx, "residual-project") + created, err := f.createMachineOnce(ctx) + require.NoError(t, err) + _, err = f.store.Organizations().DeleteProject(ctx, testOrgID, testProjectID, userPrincipal(f.userID)) + require.NoError(t, err) + var state string + require.NoError(t, + f.pool.QueryRow(ctx, + `SELECT lifecycle_state FROM machines WHERE id=$1`, + created.Machine.Machine.ID).Scan(&state)) + require.Equal(t, "deleting", state) + project := uuid.MustParse("ffffffff-ffff-4fff-8fff-ffffffffffff") + _, + err = f.pool.Exec(ctx, + `INSERT INTO projects(id,org_id,name,created_at,updated_at) VALUES($1,$2,'remaining',now(),now())`, + project, + testOrgID) + require.NoError(t, err) + var config uuid.UUID + require.NoError(t, + f.pool.QueryRow(ctx, + `INSERT INTO agent_configs(org_id,project_id,configured_model_id,source,source_format,source_hash,compiled_definition,effective_definition_hash,created_at) + SELECT org_id,$1,configured_model_id,source,source_format,source_hash,compiled_definition, + effective_definition_hash,now() + FROM agent_configs WHERE id=$2 RETURNING id`, + project, + f.agent.CurrentConfigID).Scan(&config)) + _, + err = f.store.Execution().CreateAgentFixture(ctx, + executionstore.AgentFixtureInput{ProjectID: project, + CurrentConfigID: config}) + require.NoError(t, err) + _, err = f.store.Organizations().DeleteOrganization(ctx, testOrgID, userPrincipal(f.userID)) + require.NoError(t, err) +} + +func TestGrantRevocationIncludesStartingProcessWithReleasedBinding(t *testing.T) { + ctx := t.Context() + f := newProcessDaemonFixture(t, ctx, "released-starting") + tool := createToolCallForProcessTest(t, ctx, f, "starting", "run_command") + process, err := startProcessForTest(ctx, f.Store, executionstore.ExecuteToolCallInput{ + ProjectID: testProjectID, + AgentID: f.AgentID, + RuntimeLockID: f.Lock.ID, + ToolCallID: tool}, + executionstore.CreateProcessInput{ + ExecutionSpec: processcmd.ForShell("sleep 1", "sh", ""), AgentMachineBindingID: f.BindingID, Cwd: "/work"}) + require.NoError(t, err) + accepted, found, err := acceptDaemonProcessForTest(ctx, f.Store, testOrgID, f.MachineID, f.RuntimeID, uuid.Nil) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, process.ID, accepted.Process.ID) + n, + err := f.Store.q.ReleaseExplicitAgentMachineBinding(ctx, + dbsqlc.ReleaseExplicitAgentMachineBindingParams{ProjectID: testProjectID, + AgentID: f.AgentID, + MachineID: f.MachineID}) + require.NoError(t, err) + require.EqualValues(t, 1, n) + sibling := uuid.MustParse("ffffffff-ffff-4fff-8fff-ffffffffffff") + _, + err = f.Store.pool.Exec(ctx, + `INSERT INTO agents(id,org_id,project_id,root_agent_id,state,current_config_id,created_at,updated_at) + SELECT $1,org_id,project_id,$1,'active',current_config_id,now(),now() FROM agents WHERE id=$2`, sibling, f.AgentID) + require.NoError(t, err) + _, + err = executionstore.IntegrationInsertAgentMachineBindingTx(ctx, + f.Store.q, + executionstore.IntegrationInsertAgentMachineBindingInput{ProjectID: testProjectID, + AgentID: sibling, + ProjectMachineGrantID: f.GrantID, + BindingKind: "explicit", + Cwd: "/work"}) + require.NoError(t, err) + _, err = f.Store.Execution().DeleteProjectMachineGrant(ctx, testOrgID, testProjectID, f.GrantID) + require.NoError(t, err) + var state string + require.NoError(t, f.Store.pool.QueryRow(ctx, `SELECT state FROM processes WHERE id=$1`, process.ID).Scan(&state)) + require.Equal(t, "unknown", state) +} diff --git a/internal/storage/executionstore/machine_daemon_bootstrap_store.go b/internal/storage/executionstore/machine_daemon_bootstrap_store.go index 8a7a5c434..6f2ea503b 100644 --- a/internal/storage/executionstore/machine_daemon_bootstrap_store.go +++ b/internal/storage/executionstore/machine_daemon_bootstrap_store.go @@ -333,20 +333,24 @@ func (s *Store) RevokeBYOMachineDaemonToken( orgID, machineID, tokenID uuid.UUID, reason string, ) (MachineDaemonTokenRecord, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return MachineDaemonTokenRecord{}, fmt.Errorf("begin revoke machine daemon token: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() + return agentexecution.Retry(ctx, + "RevokeBYOMachineDaemonToken", + func(ctx context.Context) (MachineDaemonTokenRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return MachineDaemonTokenRecord{}, fmt.Errorf("begin revoke machine daemon token: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() - record, err := s.RevokeBYOMachineDaemonTokenInUnit(ctx, unit, orgID, machineID, tokenID, reason) - if err != nil { - return MachineDaemonTokenRecord{}, err - } - if err := unit.Commit(ctx, "revoke machine daemon token"); err != nil { - return MachineDaemonTokenRecord{}, err - } - return record, nil + record, err := s.RevokeBYOMachineDaemonTokenInUnit(ctx, unit, orgID, machineID, tokenID, reason) + if err != nil { + return MachineDaemonTokenRecord{}, err + } + if err := unit.Commit(ctx, "revoke machine daemon token"); err != nil { + return MachineDaemonTokenRecord{}, err + } + return record, nil + }) } func (s *Store) RevokeBYOMachineDaemonTokenInUnit( @@ -544,43 +548,49 @@ func (s *Store) RecordMachineFailureReport( MachineID: input.MachineID, DaemonTokenID: input.DaemonTokenID, } - unit, err := s.cell.Begin(ctx) - if err != nil { - return fmt.Errorf("begin machine failure report: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + _, retryErr := agentexecution.Retry(ctx, + "RecordMachineFailureReport", + func(ctx context.Context) (struct{}, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return struct{}{}, fmt.Errorf("begin machine failure report: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() - qtx := dbsqlc.New(tx) - if input.Stage == MachineFailureStageDaemonUninstalled { - if _, err := qtx.LockMachineForLifecycle( - ctx, - dbsqlc.LockMachineForLifecycleParams{OrgID: input.OrgID, ID: input.MachineID}, - ); err != nil { + qtx := dbsqlc.New(tx) + if input.Stage == MachineFailureStageDaemonUninstalled { + if _, err := qtx.LockMachineForLifecycle( + ctx, + dbsqlc.LockMachineForLifecycleParams{OrgID: input.OrgID, ID: input.MachineID}, + ); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return struct{}{}, storeerr.ErrUnauthorized + } + return struct{}{}, fmt.Errorf("lock machine for daemon uninstall report: %w", err) + } + } + _, err = qtx.RecordMachineFailureReport(ctx, params) if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrUnauthorized + return struct{}{}, storeerr.ErrUnauthorized } - return fmt.Errorf("lock machine for daemon uninstall report: %w", err) - } - } - _, err = qtx.RecordMachineFailureReport(ctx, params) - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrUnauthorized - } - if err != nil { - return fmt.Errorf("record machine failure report: %w", err) - } - if input.Stage == MachineFailureStageDaemonUninstalled { - if err := completeMachineLifecycleTerminalWorkTx( - ctx, - unit, - qtx, - input.OrgID, - input.MachineID, - MachineFailureStageDaemonUninstalled, - ); err != nil { - return err - } - } - return unit.Commit(ctx, "record machine failure report") + if err != nil { + return struct{}{}, fmt.Errorf("record machine failure report: %w", err) + } + if input.Stage == MachineFailureStageDaemonUninstalled { + if err := completeMachineLifecycleTerminalWorkTx( + ctx, + unit, + qtx, + input.OrgID, + input.MachineID, + MachineFailureStageDaemonUninstalled, + ); err != nil { + return struct{}{}, err + } + } + return struct{}{}, unit.Commit(ctx, "record machine failure report") + }) + return retryErr + } diff --git a/internal/storage/executionstore/machine_daemon_reconciliation_store.go b/internal/storage/executionstore/machine_daemon_reconciliation_store.go index 291f82d05..066010c34 100644 --- a/internal/storage/executionstore/machine_daemon_reconciliation_store.go +++ b/internal/storage/executionstore/machine_daemon_reconciliation_store.go @@ -264,7 +264,7 @@ func lockAndLoadReconciliationProcesses( AgentID: record.AgentID, }) } - if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, nil); err != nil { return nil, err } return loadReconciliationProcesses( @@ -798,7 +798,7 @@ func lockAgentForProcessRecordTx(ctx context.Context, unit *agentexecution.Unit, ProjectID: record.ProjectID, ID: record.AgentID, }, - agentexecution.ExternalAuthority{}); err != nil { + nil); err != nil { return fmt.Errorf("lock agent for process authority mutation: %w", err) } return nil diff --git a/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go b/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go index 81b5c82a3..66a658b22 100644 --- a/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go +++ b/internal/storage/executionstore/machine_daemon_runtime_expiry_store.go @@ -7,6 +7,7 @@ import ( "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/notifications" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) @@ -41,48 +42,57 @@ func (s *Store) endExpiredDaemonRuntime( ctx context.Context, candidate dbsqlc.ListExpiredDaemonRuntimeCandidatesRow, ) (DaemonRuntimeRecord, bool, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonRuntimeRecord{}, false, fmt.Errorf("begin end expired daemon runtime: %w", err) + type transactionResult struct { + runtime DaemonRuntimeRecord + ended bool } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() + retried, retryErr := agentexecution.Retry(ctx, + "endExpiredDaemonRuntime", + func(ctx context.Context) (transactionResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return transactionResult{}, fmt.Errorf("begin end expired daemon runtime: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockMachineForLifecycle( - ctx, - dbsqlc.LockMachineForLifecycleParams{ - OrgID: candidate.OrgID, - ID: candidate.MachineID, - }, - ); err != nil { - return DaemonRuntimeRecord{}, false, fmt.Errorf( - "lock machine for expired daemon runtime: %w", - err, - ) - } - row, err := qtx.EndExpiredDaemonRuntime( - ctx, - dbsqlc.EndExpiredDaemonRuntimeParams(candidate), - ) - if errors.Is(err, pgx.ErrNoRows) { - if err := unit.Commit(ctx, "skipped expired daemon runtime"); err != nil { - return DaemonRuntimeRecord{}, false, err - } - return DaemonRuntimeRecord{}, false, nil - } - if err != nil { - return DaemonRuntimeRecord{}, false, fmt.Errorf("end expired daemon runtime: %w", err) - } - record := daemonRuntimeFromExpired(row) - txNotifications.AddDaemonRuntimeEnded( - record.ID, - record.MachineID, - notifications.DaemonRuntimeEndReconnect, - ) - if err := unit.Commit(ctx, "end expired daemon runtime"); err != nil { - return DaemonRuntimeRecord{}, false, err - } - return record, true, nil + qtx := dbsqlc.New(tx) + if _, err := qtx.LockMachineForLifecycle( + ctx, + dbsqlc.LockMachineForLifecycleParams{ + OrgID: candidate.OrgID, + ID: candidate.MachineID, + }, + ); err != nil { + return transactionResult{}, fmt.Errorf( + "lock machine for expired daemon runtime: %w", + err, + ) + } + row, err := qtx.EndExpiredDaemonRuntime( + ctx, + dbsqlc.EndExpiredDaemonRuntimeParams(candidate), + ) + if errors.Is(err, pgx.ErrNoRows) { + if err := unit.Commit(ctx, "skipped expired daemon runtime"); err != nil { + return transactionResult{}, err + } + return transactionResult{}, nil + } + if err != nil { + return transactionResult{}, fmt.Errorf("end expired daemon runtime: %w", err) + } + record := daemonRuntimeFromExpired(row) + txNotifications.AddDaemonRuntimeEnded( + record.ID, + record.MachineID, + notifications.DaemonRuntimeEndReconnect, + ) + if err := unit.Commit(ctx, "end expired daemon runtime"); err != nil { + return transactionResult{}, err + } + return transactionResult{record, true}, nil + }) + return retried.runtime, retried.ended, retryErr } diff --git a/internal/storage/executionstore/machine_daemon_runtime_store.go b/internal/storage/executionstore/machine_daemon_runtime_store.go index 69aa9df6f..5a8599b0a 100644 --- a/internal/storage/executionstore/machine_daemon_runtime_store.go +++ b/internal/storage/executionstore/machine_daemon_runtime_store.go @@ -12,6 +12,7 @@ import ( "github.com/omnara-ai/omnara/internal/daemonversion" "github.com/omnara-ai/omnara/internal/notifications" "github.com/omnara-ai/omnara/internal/resourcemeta" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" @@ -56,184 +57,188 @@ func (s *Store) RegisterDaemonRuntimeWithReconciliation( } input.Capacity = normalizedJSON(input.Capacity) input.ObservedPlatform = normalizedJSON(input.ObservedPlatform) - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("begin register daemon runtime: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - machine, err := qtx.LockMachineForRuntimeRegistration( - ctx, - dbsqlc.LockMachineForRuntimeRegistrationParams{ - OrgID: input.OrgID, - ID: input.MachineID, - }, - ) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("lock machine for runtime registration: %w", err) - } - instanceRow, err := qtx.GetDaemonRuntimeInstanceForUpdate( - ctx, - dbsqlc.GetDaemonRuntimeInstanceForUpdateParams{ - OrgID: input.OrgID, - MachineID: input.MachineID, - DaemonInstanceID: input.DaemonInstanceID, - }, - ) - instanceFound := err == nil - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("get daemon runtime instance: %w", err) - } - if instanceFound && (storeutil.IDFromPtr(machine.CurrentDaemonRuntimeID) != instanceRow.ID || - instanceRow.DaemonVersion != input.DaemonVersion) { - return DaemonRuntimeRegistrationRecord{}, storeerr.ErrDaemonInstanceSuperseded - } - if err := qtx.ClearMachineSleep( - ctx, - dbsqlc.ClearMachineSleepParams{OrgID: input.OrgID, ID: input.MachineID}, - ); err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("clear machine sleep on registration: %w", err) - } - priorRows, err := qtx.ListActiveDaemonRuntimesForUpdate( - ctx, - dbsqlc.ListActiveDaemonRuntimesForUpdateParams{OrgID: input.OrgID, MachineID: input.MachineID}, - ) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("list registered daemon runtimes: %w", err) - } - var runtimeID uuid.UUID - var replacedRows []dbsqlc.ListActiveDaemonRuntimesForUpdateRow - if instanceFound { - if DaemonRuntimeState(instanceRow.State) != DaemonRuntimeStateActive { - if (instanceRow.StateReasonCode != daemonRuntimeLeaseExpiredReason && - instanceRow.StateReasonCode != daemonRuntimeAsleepReason) || len(priorRows) != 0 { + return agentexecution.Retry(ctx, + "RegisterDaemonRuntimeWithReconciliation", + func(ctx context.Context) (DaemonRuntimeRegistrationRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("begin register daemon runtime: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + machine, err := qtx.LockMachineForRuntimeRegistration( + ctx, + dbsqlc.LockMachineForRuntimeRegistrationParams{ + OrgID: input.OrgID, + ID: input.MachineID, + }, + ) + if err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("lock machine for runtime registration: %w", err) + } + instanceRow, err := qtx.GetDaemonRuntimeInstanceForUpdate( + ctx, + dbsqlc.GetDaemonRuntimeInstanceForUpdateParams{ + OrgID: input.OrgID, + MachineID: input.MachineID, + DaemonInstanceID: input.DaemonInstanceID, + }, + ) + instanceFound := err == nil + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("get daemon runtime instance: %w", err) + } + if instanceFound && (storeutil.IDFromPtr(machine.CurrentDaemonRuntimeID) != instanceRow.ID || + instanceRow.DaemonVersion != input.DaemonVersion) { return DaemonRuntimeRegistrationRecord{}, storeerr.ErrDaemonInstanceSuperseded } - } - runtimeID, err = qtx.RefreshDaemonRuntimeRegistration( - ctx, - dbsqlc.RefreshDaemonRuntimeRegistrationParams{ - OrgID: input.OrgID, - MachineID: input.MachineID, - DaemonTokenID: input.DaemonTokenID, - DaemonInstanceID: input.DaemonInstanceID, - Capacity: input.Capacity, - Metadata: json.RawMessage(`{}`), - LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), - }, - ) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("refresh daemon runtime registration: %w", err) - } - } else { - replacedRows = priorRows - for _, prior := range replacedRows { - if _, err := qtx.ForceEndDaemonRuntime( + if err := qtx.ClearMachineSleep( + ctx, + dbsqlc.ClearMachineSleepParams{OrgID: input.OrgID, ID: input.MachineID}, + ); err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("clear machine sleep on registration: %w", err) + } + priorRows, err := qtx.ListActiveDaemonRuntimesForUpdate( + ctx, + dbsqlc.ListActiveDaemonRuntimesForUpdateParams{OrgID: input.OrgID, MachineID: input.MachineID}, + ) + if err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("list registered daemon runtimes: %w", err) + } + var runtimeID uuid.UUID + var replacedRows []dbsqlc.ListActiveDaemonRuntimesForUpdateRow + if instanceFound { + if DaemonRuntimeState(instanceRow.State) != DaemonRuntimeStateActive { + if (instanceRow.StateReasonCode != daemonRuntimeLeaseExpiredReason && + instanceRow.StateReasonCode != daemonRuntimeAsleepReason) || len(priorRows) != 0 { + return DaemonRuntimeRegistrationRecord{}, storeerr.ErrDaemonInstanceSuperseded + } + } + runtimeID, err = qtx.RefreshDaemonRuntimeRegistration( + ctx, + dbsqlc.RefreshDaemonRuntimeRegistrationParams{ + OrgID: input.OrgID, + MachineID: input.MachineID, + DaemonTokenID: input.DaemonTokenID, + DaemonInstanceID: input.DaemonInstanceID, + Capacity: input.Capacity, + Metadata: json.RawMessage(`{}`), + LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), + }, + ) + if err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("refresh daemon runtime registration: %w", err) + } + } else { + replacedRows = priorRows + for _, prior := range replacedRows { + if _, err := qtx.ForceEndDaemonRuntime( + ctx, + dbsqlc.ForceEndDaemonRuntimeParams{ + OrgID: prior.OrgID, + MachineID: prior.MachineID, + ID: prior.ID, + Reason: storeutil.TextFromEmpty( + daemonRuntimeReplacedReason, + ), + Message: "", + }, + ); err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("end replaced daemon runtime: %w", err) + } + txNotifications.AddDaemonRuntimeEnded( + prior.ID, + prior.MachineID, + notifications.DaemonRuntimeEndReconnect, + ) + } + runtimeID, err = qtx.InsertDaemonRuntime( + ctx, + dbsqlc.InsertDaemonRuntimeParams{ + OrgID: input.OrgID, + MachineID: input.MachineID, + DaemonTokenID: input.DaemonTokenID, + DaemonInstanceID: input.DaemonInstanceID, + DaemonVersion: input.DaemonVersion, + Capacity: input.Capacity, + Metadata: json.RawMessage(`{}`), + LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), + }, + ) + if err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("insert daemon runtime: %w", err) + } + } + if _, err := qtx.SetMachineCurrentDaemonRuntime( ctx, - dbsqlc.ForceEndDaemonRuntimeParams{ - OrgID: prior.OrgID, - MachineID: prior.MachineID, - ID: prior.ID, - Reason: storeutil.TextFromEmpty( - daemonRuntimeReplacedReason, - ), - Message: "", + dbsqlc.SetMachineCurrentDaemonRuntimeParams{ + OrgID: input.OrgID, + MachineID: input.MachineID, + DaemonRuntimeID: runtimeID, }, ); err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("end replaced daemon runtime: %w", err) + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("set current daemon runtime: %w", err) } - txNotifications.AddDaemonRuntimeEnded( - prior.ID, - prior.MachineID, - notifications.DaemonRuntimeEndReconnect, + if err := qtx.RevokeSiblingSystemMachineDaemonTokens( + ctx, + dbsqlc.RevokeSiblingSystemMachineDaemonTokensParams{ + OrgID: input.OrgID, + MachineID: input.MachineID, + ActiveTokenID: input.DaemonTokenID, + }, + ); err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("revoke sibling system bootstrap tokens: %w", err) + } + if _, err := qtx.ClearResolvedMachineFailureReport( + ctx, + dbsqlc.ClearResolvedMachineFailureReportParams{ + OrgID: input.OrgID, + MachineID: input.MachineID, + DaemonVersion: input.DaemonVersion, + }, + ); err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("clear resolved machine failure report: %w", err) + } + reconciliation, err := reconcileRegisteredRuntimeTx(ctx, unit, qtx, input) + if err != nil { + return DaemonRuntimeRegistrationRecord{}, err + } + if err := qtx.UpdateMachineObservation( + ctx, + dbsqlc.UpdateMachineObservationParams{ + OrgID: input.OrgID, + ID: input.MachineID, + ObservedPlatform: observedPlatformString(input.ObservedPlatform), + }, + ); err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("update machine observation: %w", err) + } + finalRow, err := qtx.HeartbeatDaemonRuntime( + ctx, + dbsqlc.HeartbeatDaemonRuntimeParams{ + ID: runtimeID, + OrgID: input.OrgID, + MachineID: input.MachineID, + DaemonTokenID: input.DaemonTokenID, + DaemonInstanceID: input.DaemonInstanceID, + Capacity: input.Capacity, + Metadata: json.RawMessage(`{}`), + LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), + }, ) - } - runtimeID, err = qtx.InsertDaemonRuntime( - ctx, - dbsqlc.InsertDaemonRuntimeParams{ - OrgID: input.OrgID, - MachineID: input.MachineID, - DaemonTokenID: input.DaemonTokenID, - DaemonInstanceID: input.DaemonInstanceID, - DaemonVersion: input.DaemonVersion, - Capacity: input.Capacity, - Metadata: json.RawMessage(`{}`), - LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), - }, - ) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("insert daemon runtime: %w", err) - } - } - if _, err := qtx.SetMachineCurrentDaemonRuntime( - ctx, - dbsqlc.SetMachineCurrentDaemonRuntimeParams{ - OrgID: input.OrgID, - MachineID: input.MachineID, - DaemonRuntimeID: runtimeID, - }, - ); err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("set current daemon runtime: %w", err) - } - if err := qtx.RevokeSiblingSystemMachineDaemonTokens( - ctx, - dbsqlc.RevokeSiblingSystemMachineDaemonTokensParams{ - OrgID: input.OrgID, - MachineID: input.MachineID, - ActiveTokenID: input.DaemonTokenID, - }, - ); err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("revoke sibling system bootstrap tokens: %w", err) - } - if _, err := qtx.ClearResolvedMachineFailureReport( - ctx, - dbsqlc.ClearResolvedMachineFailureReportParams{ - OrgID: input.OrgID, - MachineID: input.MachineID, - DaemonVersion: input.DaemonVersion, - }, - ); err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("clear resolved machine failure report: %w", err) - } - reconciliation, err := reconcileRegisteredRuntimeTx(ctx, unit, qtx, input) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, err - } - if err := qtx.UpdateMachineObservation( - ctx, - dbsqlc.UpdateMachineObservationParams{ - OrgID: input.OrgID, - ID: input.MachineID, - ObservedPlatform: observedPlatformString(input.ObservedPlatform), - }, - ); err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("update machine observation: %w", err) - } - finalRow, err := qtx.HeartbeatDaemonRuntime( - ctx, - dbsqlc.HeartbeatDaemonRuntimeParams{ - ID: runtimeID, - OrgID: input.OrgID, - MachineID: input.MachineID, - DaemonTokenID: input.DaemonTokenID, - DaemonInstanceID: input.DaemonInstanceID, - Capacity: input.Capacity, - Metadata: json.RawMessage(`{}`), - LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), - }, - ) - if err != nil { - return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("finalize daemon runtime registration lease: %w", err) - } - if err := unit.Commit(ctx, "register daemon runtime"); err != nil { - return DaemonRuntimeRegistrationRecord{}, err - } - return DaemonRuntimeRegistrationRecord{ - Runtime: daemonRuntimeFromHeartbeat(finalRow), - Reconciliation: reconciliation, - }, nil + if err != nil { + return DaemonRuntimeRegistrationRecord{}, fmt.Errorf("finalize daemon runtime registration lease: %w", err) + } + if err := unit.Commit(ctx, "register daemon runtime"); err != nil { + return DaemonRuntimeRegistrationRecord{}, err + } + return DaemonRuntimeRegistrationRecord{ + Runtime: daemonRuntimeFromHeartbeat(finalRow), + Reconciliation: reconciliation, + }, nil + }) } func observedPlatformString(raw json.RawMessage) string { @@ -273,55 +278,59 @@ func (s *Store) HeartbeatDaemonRuntime( return DaemonRuntimeRecord{}, errors.New("daemon runtime lease timeout is too short") } input.ObservedPlatform = normalizedJSON(input.ObservedPlatform) - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("begin heartbeat daemon runtime: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockMachineForRuntimeRegistration( - ctx, - dbsqlc.LockMachineForRuntimeRegistrationParams{ - OrgID: input.Authority.OrgID, - ID: input.Authority.MachineID, - }, - ); err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("lock machine for daemon runtime heartbeat: %w", err) - } - row, err := qtx.HeartbeatDaemonRuntime( - ctx, - dbsqlc.HeartbeatDaemonRuntimeParams{ - ID: input.Authority.DaemonRuntimeID, - OrgID: input.Authority.OrgID, - MachineID: input.Authority.MachineID, - DaemonTokenID: input.Authority.DaemonTokenID, - DaemonInstanceID: input.DaemonInstanceID, - Capacity: normalizedJSON(input.Capacity), - Metadata: json.RawMessage(`{}`), - LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered - } - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("heartbeat daemon runtime: %w", err) - } - if err := qtx.UpdateMachineObservation( - ctx, - dbsqlc.UpdateMachineObservationParams{ - OrgID: input.Authority.OrgID, - ID: input.Authority.MachineID, - ObservedPlatform: observedPlatformString(input.ObservedPlatform), - }, - ); err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("update machine heartbeat observation: %w", err) - } - if err := unit.Commit(ctx, "heartbeat daemon runtime"); err != nil { - return DaemonRuntimeRecord{}, err - } - return daemonRuntimeFromHeartbeat(row), nil + return agentexecution.Retry(ctx, + "HeartbeatDaemonRuntime", + func(ctx context.Context) (DaemonRuntimeRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("begin heartbeat daemon runtime: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) + if _, err := qtx.LockMachineForRuntimeRegistration( + ctx, + dbsqlc.LockMachineForRuntimeRegistrationParams{ + OrgID: input.Authority.OrgID, + ID: input.Authority.MachineID, + }, + ); err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("lock machine for daemon runtime heartbeat: %w", err) + } + row, err := qtx.HeartbeatDaemonRuntime( + ctx, + dbsqlc.HeartbeatDaemonRuntimeParams{ + ID: input.Authority.DaemonRuntimeID, + OrgID: input.Authority.OrgID, + MachineID: input.Authority.MachineID, + DaemonTokenID: input.Authority.DaemonTokenID, + DaemonInstanceID: input.DaemonInstanceID, + Capacity: normalizedJSON(input.Capacity), + Metadata: json.RawMessage(`{}`), + LeaseTimeoutMilliseconds: input.LeaseTimeout.Milliseconds(), + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered + } + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("heartbeat daemon runtime: %w", err) + } + if err := qtx.UpdateMachineObservation( + ctx, + dbsqlc.UpdateMachineObservationParams{ + OrgID: input.Authority.OrgID, + ID: input.Authority.MachineID, + ObservedPlatform: observedPlatformString(input.ObservedPlatform), + }, + ); err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("update machine heartbeat observation: %w", err) + } + if err := unit.Commit(ctx, "heartbeat daemon runtime"); err != nil { + return DaemonRuntimeRecord{}, err + } + return daemonRuntimeFromHeartbeat(row), nil + }) } func (s *Store) EndDaemonRuntime( @@ -331,49 +340,53 @@ func (s *Store) EndDaemonRuntime( if err := validateDaemonRuntimeAuthority(authority); err != nil { return DaemonRuntimeRecord{}, err } - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("begin end daemon runtime: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockMachineForRuntimeRegistration( - ctx, - dbsqlc.LockMachineForRuntimeRegistrationParams{ - OrgID: authority.OrgID, - ID: authority.MachineID, - }, - ); err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("lock machine for daemon runtime end: %w", err) - } - row, err := qtx.EndDaemonRuntime( - ctx, - dbsqlc.EndDaemonRuntimeParams{ - ID: authority.DaemonRuntimeID, - OrgID: authority.OrgID, - MachineID: authority.MachineID, - DaemonTokenID: authority.DaemonTokenID, - Reason: storeutil.TextFromEmpty(DaemonRuntimeReleasedReason), - Message: "", - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered - } - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("end daemon runtime: %w", err) - } - txNotifications.AddDaemonRuntimeEnded( - authority.DaemonRuntimeID, - authority.MachineID, - notifications.DaemonRuntimeEndReconnect, - ) - if err := unit.Commit(ctx, "end daemon runtime"); err != nil { - return DaemonRuntimeRecord{}, err - } - return daemonRuntimeFromEnd(row), nil + return agentexecution.Retry(ctx, + "EndDaemonRuntime", + func(ctx context.Context) (DaemonRuntimeRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("begin end daemon runtime: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + if _, err := qtx.LockMachineForRuntimeRegistration( + ctx, + dbsqlc.LockMachineForRuntimeRegistrationParams{ + OrgID: authority.OrgID, + ID: authority.MachineID, + }, + ); err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("lock machine for daemon runtime end: %w", err) + } + row, err := qtx.EndDaemonRuntime( + ctx, + dbsqlc.EndDaemonRuntimeParams{ + ID: authority.DaemonRuntimeID, + OrgID: authority.OrgID, + MachineID: authority.MachineID, + DaemonTokenID: authority.DaemonTokenID, + Reason: storeutil.TextFromEmpty(DaemonRuntimeReleasedReason), + Message: "", + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered + } + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("end daemon runtime: %w", err) + } + txNotifications.AddDaemonRuntimeEnded( + authority.DaemonRuntimeID, + authority.MachineID, + notifications.DaemonRuntimeEndReconnect, + ) + if err := unit.Commit(ctx, "end daemon runtime"); err != nil { + return DaemonRuntimeRecord{}, err + } + return daemonRuntimeFromEnd(row), nil + }) } func (s *Store) SleepDaemonRuntime( @@ -383,74 +396,78 @@ func (s *Store) SleepDaemonRuntime( if err := validateDaemonRuntimeAuthority(authority); err != nil { return DaemonRuntimeRecord{}, err } - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("begin sleep daemon runtime: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockMachineForRuntimeRegistration( - ctx, - dbsqlc.LockMachineForRuntimeRegistrationParams{ - OrgID: authority.OrgID, - ID: authority.MachineID, - }, - ); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered - } - return DaemonRuntimeRecord{}, fmt.Errorf("lock machine for daemon runtime sleep: %w", err) - } - unfinished, err := qtx.MachineHasUnfinishedDaemonWork( - ctx, - dbsqlc.MachineHasUnfinishedDaemonWorkParams{ - OrgID: authority.OrgID, - MachineID: authority.MachineID, - }, - ) - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("check unfinished daemon work for sleep: %w", err) - } - if unfinished != nil && *unfinished { - return DaemonRuntimeRecord{}, storeerr.ErrMachineSleepPendingWork - } - if _, err := qtx.MarkMachineAsleep( - ctx, - dbsqlc.MarkMachineAsleepParams{OrgID: authority.OrgID, ID: authority.MachineID}, - ); err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return DaemonRuntimeRecord{}, storeerr.ErrMachineNotWakeCapable - } - return DaemonRuntimeRecord{}, fmt.Errorf("mark machine asleep: %w", err) - } - row, err := qtx.EndDaemonRuntime( - ctx, - dbsqlc.EndDaemonRuntimeParams{ - ID: authority.DaemonRuntimeID, - OrgID: authority.OrgID, - MachineID: authority.MachineID, - DaemonTokenID: authority.DaemonTokenID, - Reason: storeutil.TextFromEmpty(daemonRuntimeAsleepReason), - Message: "", - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered - } - if err != nil { - return DaemonRuntimeRecord{}, fmt.Errorf("end daemon runtime for sleep: %w", err) - } - txNotifications.AddDaemonRuntimeEnded( - authority.DaemonRuntimeID, - authority.MachineID, - notifications.DaemonRuntimeEndReconnect, - ) - if err := unit.Commit(ctx, "sleep daemon runtime"); err != nil { - return DaemonRuntimeRecord{}, err - } - return daemonRuntimeFromEnd(row), nil + return agentexecution.Retry(ctx, + "SleepDaemonRuntime", + func(ctx context.Context) (DaemonRuntimeRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("begin sleep daemon runtime: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + if _, err := qtx.LockMachineForRuntimeRegistration( + ctx, + dbsqlc.LockMachineForRuntimeRegistrationParams{ + OrgID: authority.OrgID, + ID: authority.MachineID, + }, + ); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered + } + return DaemonRuntimeRecord{}, fmt.Errorf("lock machine for daemon runtime sleep: %w", err) + } + unfinished, err := qtx.MachineHasUnfinishedDaemonWork( + ctx, + dbsqlc.MachineHasUnfinishedDaemonWorkParams{ + OrgID: authority.OrgID, + MachineID: authority.MachineID, + }, + ) + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("check unfinished daemon work for sleep: %w", err) + } + if unfinished != nil && *unfinished { + return DaemonRuntimeRecord{}, storeerr.ErrMachineSleepPendingWork + } + if _, err := qtx.MarkMachineAsleep( + ctx, + dbsqlc.MarkMachineAsleepParams{OrgID: authority.OrgID, ID: authority.MachineID}, + ); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return DaemonRuntimeRecord{}, storeerr.ErrMachineNotWakeCapable + } + return DaemonRuntimeRecord{}, fmt.Errorf("mark machine asleep: %w", err) + } + row, err := qtx.EndDaemonRuntime( + ctx, + dbsqlc.EndDaemonRuntimeParams{ + ID: authority.DaemonRuntimeID, + OrgID: authority.OrgID, + MachineID: authority.MachineID, + DaemonTokenID: authority.DaemonTokenID, + Reason: storeutil.TextFromEmpty(daemonRuntimeAsleepReason), + Message: "", + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return DaemonRuntimeRecord{}, storeerr.ErrDaemonRuntimeUnregistered + } + if err != nil { + return DaemonRuntimeRecord{}, fmt.Errorf("end daemon runtime for sleep: %w", err) + } + txNotifications.AddDaemonRuntimeEnded( + authority.DaemonRuntimeID, + authority.MachineID, + notifications.DaemonRuntimeEndReconnect, + ) + if err := unit.Commit(ctx, "sleep daemon runtime"); err != nil { + return DaemonRuntimeRecord{}, err + } + return daemonRuntimeFromEnd(row), nil + }) } func (s *Store) RegisteredDaemonRuntimeExists(ctx context.Context, authority DaemonRuntimeAuthority) (bool, error) { diff --git a/internal/storage/executionstore/machine_daemon_store.go b/internal/storage/executionstore/machine_daemon_store.go index 46da24739..51a312d1f 100644 --- a/internal/storage/executionstore/machine_daemon_store.go +++ b/internal/storage/executionstore/machine_daemon_store.go @@ -353,9 +353,11 @@ func (s *Store) UpdateMachine(ctx context.Context, input UpdateMachineInput) (Ma if input.OrgID == uuid.Nil || input.MachineID == uuid.Nil { return MachineRecord{}, errors.New("org and machine are required") } - return storeutil.RetryTransaction(ctx, "update_machine", func() (MachineRecord, error) { - return s.updateMachineOnce(ctx, input) - }) + return agentexecution.Retry(ctx, + "update_machine", + func(ctx context.Context) (MachineRecord, error) { + return s.updateMachineOnce(ctx, input) + }) } func (s *Store) updateMachineOnce(ctx context.Context, input UpdateMachineInput) (MachineRecord, error) { @@ -544,9 +546,11 @@ func (s *Store) DeleteMachine( if input.OrgID == uuid.Nil || input.MachineID == uuid.Nil { return MachineRecord{}, errors.New("org and machine are required") } - return storeutil.RetryTransaction(ctx, "delete_machine", func() (MachineRecord, error) { - return s.deleteMachineOnce(ctx, input) - }) + return agentexecution.Retry(ctx, + "delete_machine", + func(ctx context.Context) (MachineRecord, error) { + return s.deleteMachineOnce(ctx, input) + }) } func (s *Store) deleteMachineOnce( @@ -999,9 +1003,12 @@ func (s *Store) DeleteProjectMachineGrant( ctx context.Context, orgID, projectID, id uuid.UUID, ) (ProjectMachineGrantRecord, error) { - return storeutil.RetryTransaction(ctx, "delete_project_machine_grant", func() (ProjectMachineGrantRecord, error) { - return s.deleteProjectMachineGrantOnce(ctx, orgID, projectID, id) - }) + return agentexecution.Retry(ctx, + "delete_project_machine_grant", + func(ctx context.Context) (ProjectMachineGrantRecord, + error) { + return s.deleteProjectMachineGrantOnce(ctx, orgID, projectID, id) + }) } func (s *Store) deleteProjectMachineGrantOnce( diff --git a/internal/storage/executionstore/machine_lifecycle_store.go b/internal/storage/executionstore/machine_lifecycle_store.go index a1dcaa6b4..02463e3e9 100644 --- a/internal/storage/executionstore/machine_lifecycle_store.go +++ b/internal/storage/executionstore/machine_lifecycle_store.go @@ -647,52 +647,61 @@ func (s *Store) ClaimExpiredIdlePoolMachineDeletion( if orgID == uuid.Nil || machineID == uuid.Nil { return PoolMachineDeletionClaim{}, false, errors.New("org and machine are required") } - unit, err := s.cell.Begin(ctx) - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf("begin claim expired idle pool machine deletion: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockMachineForLifecycle( - ctx, - dbsqlc.LockMachineForLifecycleParams{OrgID: orgID, ID: machineID}, - ); errors.Is(err, pgx.ErrNoRows) { - return PoolMachineDeletionClaim{}, false, nil - } else if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf("lock expired idle pool machine deletion: %w", err) - } - row, err := qtx.ClaimExpiredIdlePoolMachineDeletion( - ctx, - dbsqlc.ClaimExpiredIdlePoolMachineDeletionParams{ - OrgID: orgID, - ID: machineID, - ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return PoolMachineDeletionClaim{}, false, nil - } - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf("claim expired idle pool machine deletion: %w", err) - } - claim := expiredIdlePoolMachineDeletionClaimFromSQLC(row) - claim, err = s.finalizePoolMachineDeletionClaimTx( - ctx, - unit, - qtx, - txNotifications, - claim, - machineDeletingReason, - ) - if err != nil { - return PoolMachineDeletionClaim{}, false, err - } - if err := unit.Commit(ctx, "claim expired idle pool machine deletion"); err != nil { - return PoolMachineDeletionClaim{}, false, err - } - return claim, true, nil + type transactionResult struct { + claim PoolMachineDeletionClaim + claimed bool + } + retried, retryErr := agentexecution.Retry(ctx, + "ClaimExpiredIdlePoolMachineDeletion", + func(ctx context.Context) (transactionResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return transactionResult{}, fmt.Errorf("begin claim expired idle pool machine deletion: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + if _, err := qtx.LockMachineForLifecycle( + ctx, + dbsqlc.LockMachineForLifecycleParams{OrgID: orgID, ID: machineID}, + ); errors.Is(err, pgx.ErrNoRows) { + return transactionResult{}, nil + } else if err != nil { + return transactionResult{}, fmt.Errorf("lock expired idle pool machine deletion: %w", err) + } + row, err := qtx.ClaimExpiredIdlePoolMachineDeletion( + ctx, + dbsqlc.ClaimExpiredIdlePoolMachineDeletionParams{ + OrgID: orgID, + ID: machineID, + ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return transactionResult{}, nil + } + if err != nil { + return transactionResult{}, fmt.Errorf("claim expired idle pool machine deletion: %w", err) + } + claim := expiredIdlePoolMachineDeletionClaimFromSQLC(row) + claim, err = s.finalizePoolMachineDeletionClaimTx( + ctx, + unit, + qtx, + txNotifications, + claim, + machineDeletingReason, + ) + if err != nil { + return transactionResult{}, err + } + if err := unit.Commit(ctx, "claim expired idle pool machine deletion"); err != nil { + return transactionResult{}, err + } + return transactionResult{claim, true}, nil + }) + return retried.claim, retried.claimed, retryErr } func (s *Store) ClaimPoolMachineDeletion( @@ -708,55 +717,64 @@ func (s *Store) ClaimPoolMachineDeletion( if input.ExpectedLifecycleVersion <= 0 { return PoolMachineDeletionClaim{}, false, errors.New("expected lifecycle version is required") } - unit, err := s.cell.Begin(ctx) - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf("begin claim pool machine deletion: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - if _, err := qtx.LockMachineForLifecycle( - ctx, - dbsqlc.LockMachineForLifecycleParams{OrgID: input.OrgID, ID: input.MachineID}, - ); errors.Is(err, pgx.ErrNoRows) { - return PoolMachineDeletionClaim{}, false, nil - } else if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf("lock pool machine for deletion claim: %w", err) - } - row, err := qtx.ClaimPoolMachineDeletion(ctx, dbsqlc.ClaimPoolMachineDeletionParams{ - OrgID: input.OrgID, - ID: input.MachineID, - LifecycleReasonCode: storeutil.TextFromEmpty(input.LifecycleReasonCode), - LifecycleReasonMessage: input.LifecycleReasonMessage, - ExpectedLifecycleVersion: input.ExpectedLifecycleVersion, - MaxProvisionAttempts: DefaultPoolMachineProvisionFailureLimit, - StaleBootstrapAgeSeconds: int64(StaleMachineBootstrapAge / time.Second), - ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), - MissingResourceFinalityAgeSeconds: int64(missingProviderResourceFinalityAge / time.Second), - }) - if errors.Is(err, pgx.ErrNoRows) { - return PoolMachineDeletionClaim{}, false, nil - } - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf("claim pool machine deletion: %w", err) - } - claim := poolMachineDeletionClaimFromSQLC(row) - claim, err = s.finalizePoolMachineDeletionClaimTx( - ctx, - unit, - qtx, - txNotifications, - claim, - machineDeletingReason, - ) - if err != nil { - return PoolMachineDeletionClaim{}, false, err - } - if err := unit.Commit(ctx, "claim pool machine deletion"); err != nil { - return PoolMachineDeletionClaim{}, false, err - } - return claim, true, nil + type transactionResult struct { + claim PoolMachineDeletionClaim + claimed bool + } + retried, retryErr := agentexecution.Retry(ctx, + "ClaimPoolMachineDeletion", + func(ctx context.Context) (transactionResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return transactionResult{}, fmt.Errorf("begin claim pool machine deletion: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + if _, err := qtx.LockMachineForLifecycle( + ctx, + dbsqlc.LockMachineForLifecycleParams{OrgID: input.OrgID, ID: input.MachineID}, + ); errors.Is(err, pgx.ErrNoRows) { + return transactionResult{}, nil + } else if err != nil { + return transactionResult{}, fmt.Errorf("lock pool machine for deletion claim: %w", err) + } + row, err := qtx.ClaimPoolMachineDeletion(ctx, dbsqlc.ClaimPoolMachineDeletionParams{ + OrgID: input.OrgID, + ID: input.MachineID, + LifecycleReasonCode: storeutil.TextFromEmpty(input.LifecycleReasonCode), + LifecycleReasonMessage: input.LifecycleReasonMessage, + ExpectedLifecycleVersion: input.ExpectedLifecycleVersion, + MaxProvisionAttempts: DefaultPoolMachineProvisionFailureLimit, + StaleBootstrapAgeSeconds: int64(StaleMachineBootstrapAge / time.Second), + ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), + MissingResourceFinalityAgeSeconds: int64(missingProviderResourceFinalityAge / time.Second), + }) + if errors.Is(err, pgx.ErrNoRows) { + return transactionResult{}, nil + } + if err != nil { + return transactionResult{}, fmt.Errorf("claim pool machine deletion: %w", err) + } + claim := poolMachineDeletionClaimFromSQLC(row) + claim, err = s.finalizePoolMachineDeletionClaimTx( + ctx, + unit, + qtx, + txNotifications, + claim, + machineDeletingReason, + ) + if err != nil { + return transactionResult{}, err + } + if err := unit.Commit(ctx, "claim pool machine deletion"); err != nil { + return transactionResult{}, err + } + return transactionResult{claim, true}, nil + }) + return retried.claim, retried.claimed, retryErr } func (s *Store) finalizePoolMachineDeletionClaimTx( @@ -935,128 +953,134 @@ func (s *Store) CompletePoolMachineDeletion( if deleteAttempt <= 0 { return errors.New("delete attempt is required") } - unit, err := s.cell.Begin(ctx) - if err != nil { - return fmt.Errorf("begin complete pool machine deletion: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - poolID, err := qtx.GetPoolMachinePoolIDForLifecycle( - ctx, - dbsqlc.GetPoolMachinePoolIDForLifecycleParams{OrgID: orgID, ID: machineID}, - ) - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrStateTransitionConflict - } - if err != nil { - return fmt.Errorf("load machine pool for deletion completion: %w", err) - } - if poolID == nil { - return errors.New("pool machine has no machine pool") - } - if err := lifecyclelock.Pools( - ctx, - tx, - []lifecyclelock.PoolRef{{OrgID: orgID, PoolID: *poolID}}, - ); err != nil { - return err - } - lockedMachine, err := qtx.LockPoolMachineDeletionAttemptForLifecycle( - ctx, - dbsqlc.LockPoolMachineDeletionAttemptForLifecycleParams{ - OrgID: orgID, - ID: machineID, - DeleteAttempt: deleteAttempt, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return storeerr.ErrStateTransitionConflict - } - if err != nil { - return fmt.Errorf("lock machine for deletion completion: %w", err) - } - if lockedMachine.MachinePoolID == nil || *lockedMachine.MachinePoolID != *poolID { - return errors.New("pool machine changed machine pool during deletion completion") - } - active, err := qtx.ListActiveDaemonRuntimesForUpdate( - ctx, - dbsqlc.ListActiveDaemonRuntimesForUpdateParams{OrgID: orgID, MachineID: machineID}, - ) - if err != nil { - return fmt.Errorf("list active runtimes for deletion completion: %w", err) - } - for _, runtime := range active { - if _, err := qtx.ForceEndDaemonRuntime(ctx, dbsqlc.ForceEndDaemonRuntimeParams{ - OrgID: orgID, - MachineID: machineID, - ID: runtime.ID, - Reason: storeutil.TextFromEmpty("machine_deleted"), - Message: "", - }); err != nil && - !errors.Is(err, pgx.ErrNoRows) { - return fmt.Errorf("end runtime for deletion completion: %w", err) - } - txNotifications.AddDaemonRuntimeEnded( - runtime.ID, - runtime.MachineID, - notifications.DaemonRuntimeEndMachineDecommissioned, - ) - } - if err := completeMachineLifecycleTerminalWorkTx( - ctx, - unit, - qtx, - orgID, - machineID, - "machine_deleted", - ); err != nil { - return err - } - if err := qtx.RevokeMachineDaemonTokensForMachine(ctx, dbsqlc.RevokeMachineDaemonTokensForMachineParams{ - OrgID: orgID, - MachineID: machineID, - Reason: "machine_deleted", - }); err != nil { - return fmt.Errorf("revoke machine daemon tokens for deletion completion: %w", err) - } - if err := qtx.ReleaseAgentMachineBindingsForMachine(ctx, dbsqlc.ReleaseAgentMachineBindingsForMachineParams{ - OrgID: orgID, - MachineID: machineID, - }); err != nil { - return fmt.Errorf("release agent machine bindings for deletion completion: %w", err) - } - if err := qtx.DeletePoolProjectMachineGrantsForMachine( - ctx, - dbsqlc.DeletePoolProjectMachineGrantsForMachineParams{ - OrgID: orgID, - MachineID: machineID, - }, - ); err != nil { - return fmt.Errorf("delete pool project machine grants for deletion completion: %w", err) - } - if _, err := qtx.DeletePoolMachine(ctx, dbsqlc.DeletePoolMachineParams{ - OrgID: orgID, - ID: machineID, - DeleteAttempt: deleteAttempt, - }); err != nil { - return fmt.Errorf("delete machine for deletion completion: %w", err) - } - if err := qtx.ReleaseMachinePoolCredentialIfIdle(ctx, dbsqlc.ReleaseMachinePoolCredentialIfIdleParams{ - OrgID: orgID, - MachinePoolID: *poolID, - }); err != nil { - return fmt.Errorf("release machine pool credential for deletion completion: %w", err) - } - if _, err := qtx.DestroyUnreferencedSecretVersionsForDeletedOrg( - ctx, - dbsqlc.DestroyUnreferencedSecretVersionsForDeletedOrgParams{OrgID: orgID}, - ); err != nil { - return fmt.Errorf("destroy deleted-org secret versions for deletion completion: %w", err) - } - if err := unit.Commit(ctx, "complete pool machine deletion"); err != nil { - return err - } - return nil + _, retryErr := agentexecution.Retry(ctx, + "CompletePoolMachineDeletion", + func(ctx context.Context) (struct{}, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return struct{}{}, fmt.Errorf("begin complete pool machine deletion: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + poolID, err := qtx.GetPoolMachinePoolIDForLifecycle( + ctx, + dbsqlc.GetPoolMachinePoolIDForLifecycleParams{OrgID: orgID, ID: machineID}, + ) + if errors.Is(err, pgx.ErrNoRows) { + return struct{}{}, storeerr.ErrStateTransitionConflict + } + if err != nil { + return struct{}{}, fmt.Errorf("load machine pool for deletion completion: %w", err) + } + if poolID == nil { + return struct{}{}, errors.New("pool machine has no machine pool") + } + if err := lifecyclelock.Pools( + ctx, + tx, + []lifecyclelock.PoolRef{{OrgID: orgID, PoolID: *poolID}}, + ); err != nil { + return struct{}{}, err + } + lockedMachine, err := qtx.LockPoolMachineDeletionAttemptForLifecycle( + ctx, + dbsqlc.LockPoolMachineDeletionAttemptForLifecycleParams{ + OrgID: orgID, + ID: machineID, + DeleteAttempt: deleteAttempt, + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return struct{}{}, storeerr.ErrStateTransitionConflict + } + if err != nil { + return struct{}{}, fmt.Errorf("lock machine for deletion completion: %w", err) + } + if lockedMachine.MachinePoolID == nil || *lockedMachine.MachinePoolID != *poolID { + return struct{}{}, errors.New("pool machine changed machine pool during deletion completion") + } + active, err := qtx.ListActiveDaemonRuntimesForUpdate( + ctx, + dbsqlc.ListActiveDaemonRuntimesForUpdateParams{OrgID: orgID, MachineID: machineID}, + ) + if err != nil { + return struct{}{}, fmt.Errorf("list active runtimes for deletion completion: %w", err) + } + for _, runtime := range active { + if _, err := qtx.ForceEndDaemonRuntime(ctx, dbsqlc.ForceEndDaemonRuntimeParams{ + OrgID: orgID, + MachineID: machineID, + ID: runtime.ID, + Reason: storeutil.TextFromEmpty("machine_deleted"), + Message: "", + }); err != nil && + !errors.Is(err, pgx.ErrNoRows) { + return struct{}{}, fmt.Errorf("end runtime for deletion completion: %w", err) + } + txNotifications.AddDaemonRuntimeEnded( + runtime.ID, + runtime.MachineID, + notifications.DaemonRuntimeEndMachineDecommissioned, + ) + } + if err := completeMachineLifecycleTerminalWorkTx( + ctx, + unit, + qtx, + orgID, + machineID, + "machine_deleted", + ); err != nil { + return struct{}{}, err + } + if err := qtx.RevokeMachineDaemonTokensForMachine(ctx, dbsqlc.RevokeMachineDaemonTokensForMachineParams{ + OrgID: orgID, + MachineID: machineID, + Reason: "machine_deleted", + }); err != nil { + return struct{}{}, fmt.Errorf("revoke machine daemon tokens for deletion completion: %w", err) + } + if err := qtx.ReleaseAgentMachineBindingsForMachine(ctx, dbsqlc.ReleaseAgentMachineBindingsForMachineParams{ + OrgID: orgID, + MachineID: machineID, + }); err != nil { + return struct{}{}, fmt.Errorf("release agent machine bindings for deletion completion: %w", err) + } + if err := qtx.DeletePoolProjectMachineGrantsForMachine( + ctx, + dbsqlc.DeletePoolProjectMachineGrantsForMachineParams{ + OrgID: orgID, + MachineID: machineID, + }, + ); err != nil { + return struct{}{}, fmt.Errorf("delete pool project machine grants for deletion completion: %w", err) + } + if _, err := qtx.DeletePoolMachine(ctx, dbsqlc.DeletePoolMachineParams{ + OrgID: orgID, + ID: machineID, + DeleteAttempt: deleteAttempt, + }); err != nil { + return struct{}{}, fmt.Errorf("delete machine for deletion completion: %w", err) + } + if err := qtx.ReleaseMachinePoolCredentialIfIdle(ctx, dbsqlc.ReleaseMachinePoolCredentialIfIdleParams{ + OrgID: orgID, + MachinePoolID: *poolID, + }); err != nil { + return struct{}{}, fmt.Errorf("release machine pool credential for deletion completion: %w", err) + } + if _, err := qtx.DestroyUnreferencedSecretVersionsForDeletedOrg( + ctx, + dbsqlc.DestroyUnreferencedSecretVersionsForDeletedOrgParams{OrgID: orgID}, + ); err != nil { + return struct{}{}, fmt.Errorf("destroy deleted-org secret versions for deletion completion: %w", err) + } + if err := unit.Commit(ctx, "complete pool machine deletion"); err != nil { + return struct{}{}, err + } + return struct{}{}, nil + }) + return retryErr + } diff --git a/internal/storage/executionstore/machine_lifecycle_terminal_work.go b/internal/storage/executionstore/machine_lifecycle_terminal_work.go index 140ce5b28..6595847b0 100644 --- a/internal/storage/executionstore/machine_lifecycle_terminal_work.go +++ b/internal/storage/executionstore/machine_lifecycle_terminal_work.go @@ -37,7 +37,7 @@ func completeMachineLifecycleTerminalWorkTx( AgentID: agent.AgentID, }) } - if err := unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}); err != nil { + if err := unit.LockAgentRefs(ctx, refs, nil); err != nil { return err } if err := completeMachineLifecycleTerminalProcessesTx( @@ -75,23 +75,6 @@ func completeExecutionRevokedProcessesTx( scope executionRevokedProcessScope, reason string, ) error { - lockParams := dbsqlc.ListAgentsForExecutionRevokedParams{ - ProjectID: scope.projectID, - AgentID: storeutil.IDFromNil(scope.agentID), - ProjectMachineGrantID: storeutil.IDFromNil(scope.projectMachineGrantID), - ProjectMachinePoolGrantID: storeutil.IDFromNil(scope.projectMachinePoolGrantID), - } - agents, err := qtx.ListAgentsForExecutionRevoked(ctx, lockParams) - if err != nil { - return fmt.Errorf("list agents for execution revoke: %w", err) - } - refs := make([]lifecyclelock.AgentRef, 0, len(agents)) - for _, agentID := range agents { - refs = append(refs, lifecyclelock.AgentRef{ProjectID: scope.projectID, AgentID: agentID}) - } - if err := unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}); err != nil { - return err - } rows, err := qtx.ListProcessesForExecutionRevoked( ctx, dbsqlc.ListProcessesForExecutionRevokedParams{ @@ -104,6 +87,13 @@ func completeExecutionRevokedProcessesTx( if err != nil { return fmt.Errorf("list processes for execution revoke: %w", err) } + refs := make([]lifecyclelock.AgentRef, 0, len(rows)) + for _, row := range rows { + refs = append(refs, lifecyclelock.AgentRef{ProjectID: row.ProjectID, AgentID: row.AgentID}) + } + if err := unit.LockAgentRefs(ctx, refs, nil); err != nil { + return err + } for _, row := range rows { process := processRecordFromSQLC(row) switch process.State { diff --git a/internal/storage/executionstore/machine_pools_store.go b/internal/storage/executionstore/machine_pools_store.go index 117e2a1fb..5ba80b11d 100644 --- a/internal/storage/executionstore/machine_pools_store.go +++ b/internal/storage/executionstore/machine_pools_store.go @@ -987,9 +987,11 @@ func (s *Store) DeleteMachinePool(ctx context.Context, orgID, id uuid.UUID) ([]M if pool.ManagementKind == management.Cluster { return nil, fmt.Errorf("cluster-managed machine pools cannot be deleted: %w", storeerr.ErrStateTransitionConflict) } - return storeutil.RetryTransaction(ctx, "delete_machine_pool", func() ([]MachineRecord, error) { - return s.deleteMachinePoolOnce(ctx, orgID, id) - }) + return agentexecution.Retry(ctx, + "delete_machine_pool", + func(ctx context.Context) ([]MachineRecord, error) { + return s.deleteMachinePoolOnce(ctx, orgID, id) + }) } func (s *Store) deleteMachinePoolOnce(ctx context.Context, orgID, id uuid.UUID) ([]MachineRecord, error) { @@ -1075,7 +1077,19 @@ func (s *Store) DeleteMachinePoolInUnit( AgentID: agentRow.AgentID, }) } - if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { + for _, machineID := range machineIDs { + rows, + err := qtx.ListMachineLifecycleTerminalAgentRefs(ctx, + dbsqlc.ListMachineLifecycleTerminalAgentRefsParams{OrgID: orgID, + MachineID: machineID}) + if err != nil { + return nil, err + } + for _, row := range rows { + agentRefs = append(agentRefs, lifecyclelock.AgentRef{ProjectID: row.ProjectID, AgentID: row.AgentID}) + } + } + if err := unit.LockAgentRefs(ctx, agentRefs, nil); err != nil { return nil, err } if _, err := qtx.DeleteMachinePool( diff --git a/internal/storage/executionstore/machine_wake_store.go b/internal/storage/executionstore/machine_wake_store.go index 6ca8d1954..dd5bf8e62 100644 --- a/internal/storage/executionstore/machine_wake_store.go +++ b/internal/storage/executionstore/machine_wake_store.go @@ -8,9 +8,9 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -36,9 +36,11 @@ func (s *Store) BeginMachineWake( return MachineWakeUnavailable, errors.New("machine wake timeout must be at least one millisecond") } - return storeutil.RetryTransaction(ctx, "begin_machine_wake", func() (MachineWakeDisposition, error) { - return s.beginMachineWakeOnce(ctx, orgID, machineID, machinePoolID, wakeTimeout) - }) + return agentexecution.Retry(ctx, + "begin_machine_wake", + func(ctx context.Context) (MachineWakeDisposition, error) { + return s.beginMachineWakeOnce(ctx, orgID, machineID, machinePoolID, wakeTimeout) + }) } func (s *Store) beginMachineWakeOnce( diff --git a/internal/storage/executionstore/model_contexts_claims.go b/internal/storage/executionstore/model_contexts_claims.go index 0c88910bc..0b21068b3 100644 --- a/internal/storage/executionstore/model_contexts_claims.go +++ b/internal/storage/executionstore/model_contexts_claims.go @@ -16,59 +16,67 @@ func (s *Store) ClaimCompactionModelCall( ctx context.Context, input ClaimCompactionModelCallInput, ) (ModelCallClaim, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return ModelCallClaim{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - result, err := h.ResumeCompaction(ctx, agentexecution.ResumeCompactionInput{ - RuntimeLockID: input.RuntimeLockID, ParentID: input.ParentContextID, - Watermark: input.InputEventSequence, SourceEnd: input.SourceEventSequenceEnd, - }) - if err != nil { - return ModelCallClaim{}, err - } - if err := applyAdmissionDestination(ctx, - unit, - input.ProjectID, - input.AgentID, - result.Admission); err != nil { - return ModelCallClaim{}, err - } - if result.Preempted { - return ModelCallClaim{}, storeerr.ErrAgentNotAdvanceable - } - claim, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, result.Model) - if err != nil { - return ModelCallClaim{}, err - } - if err = unit.Commit(ctx, "claim compaction"); err != nil { - return ModelCallClaim{}, err - } - return claim, nil + return agentexecution.Retry(ctx, + "ClaimCompactionModelCall", + func(ctx context.Context) (ModelCallClaim, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return ModelCallClaim{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.ResumeCompaction(ctx, agentexecution.ResumeCompactionInput{ + RuntimeLockID: input.RuntimeLockID, ParentID: input.ParentContextID, + Watermark: input.InputEventSequence, SourceEnd: input.SourceEventSequenceEnd, + }) + if err != nil { + return ModelCallClaim{}, err + } + if err := applyAdmissionDestination(ctx, + unit, + input.ProjectID, + input.AgentID, + result.Admission); err != nil { + return ModelCallClaim{}, err + } + if result.Preempted { + return ModelCallClaim{}, storeerr.ErrAgentNotAdvanceable + } + claim, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, result.Model) + if err != nil { + return ModelCallClaim{}, err + } + if err = unit.Commit(ctx, "claim compaction"); err != nil { + return ModelCallClaim{}, err + } + return claim, nil + }) } func (s *Store) ClaimNextModelCallContext( ctx context.Context, input ClaimNextModelCallContextInput, ) (ModelCallClaim, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return ModelCallClaim{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - prepared, err := h.ResumeModel(ctx, input.RuntimeLockID, input.PredecessorModelCallContextID) - if err != nil { - return ModelCallClaim{}, err - } - result, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, prepared) - if err != nil { - return ModelCallClaim{}, err - } - if err = unit.Commit(ctx, "claim retry"); err != nil { - return ModelCallClaim{}, err - } - return result, nil + return agentexecution.Retry(ctx, + "ClaimNextModelCallContext", + func(ctx context.Context) (ModelCallClaim, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return ModelCallClaim{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + prepared, err := h.ResumeModel(ctx, input.RuntimeLockID, input.PredecessorModelCallContextID) + if err != nil { + return ModelCallClaim{}, err + } + result, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, prepared) + if err != nil { + return ModelCallClaim{}, err + } + if err = unit.Commit(ctx, "claim retry"); err != nil { + return ModelCallClaim{}, err + } + return result, nil + }) } func (s *Store) GetModelCallContext( diff --git a/internal/storage/executionstore/model_contexts_compaction.go b/internal/storage/executionstore/model_contexts_compaction.go index abf9c6cc6..d0a39345c 100644 --- a/internal/storage/executionstore/model_contexts_compaction.go +++ b/internal/storage/executionstore/model_contexts_compaction.go @@ -4,7 +4,6 @@ import ( "context" "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" - "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -13,96 +12,106 @@ func (s *Store) RecordModelCallFailureAndClaimCompaction( ctx context.Context, input RecordModelCallFailureAndClaimCompactionInput, ) (TriggeredCompactionHandoff, error) { - f := input.Failure - if input.ParentContextID != f.ModelCallContextID || f.RecoveryKind != ModelCallRecoveryCompact { - return TriggeredCompactionHandoff{}, storeerr.ErrStateTransitionConflict - } - unit, h, err := beginExecution(ctx, s, f.ProjectID, f.AgentID) - if err != nil { - return TriggeredCompactionHandoff{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - result, err := h.BeginCompaction( - ctx, - agentexecution.BeginCompactionInput{ - Failure: executionFailure(f), - SourceEnd: input.SourceEventSequenceEnd, - }, - ) - if err != nil { - return TriggeredCompactionHandoff{}, err - } - if err := applyAdmissionDestination(ctx, - unit, - input.Failure.ProjectID, - input.Failure.AgentID, - result.Admission); err != nil { - return TriggeredCompactionHandoff{}, err - } - parent, err := loadModelCallContextByID( - ctx, - dbsqlc.New(unit.DB()), - f.ProjectID, - f.AgentID, - input.ParentContextID, - ) - if err != nil { - return TriggeredCompactionHandoff{}, err - } - claim, err := executionClaim(ctx, unit, f.ProjectID, f.AgentID, result.Model) - if err != nil { - return TriggeredCompactionHandoff{}, err - } - if err = unit.Commit(ctx, "begin compaction"); err != nil { - return TriggeredCompactionHandoff{}, err - } - return TriggeredCompactionHandoff{ - ParentContext: parent, - CompactionCall: claim, - BoundaryPreempted: result.Preempted, - }, nil + return agentexecution.Retry(ctx, + "RecordModelCallFailureAndClaimCompaction", + func(ctx context.Context) (TriggeredCompactionHandoff, + error) { + f := input.Failure + if input.ParentContextID != f.ModelCallContextID || f.RecoveryKind != ModelCallRecoveryCompact { + return TriggeredCompactionHandoff{}, storeerr.ErrStateTransitionConflict + } + unit, h, err := beginExecution(ctx, s, f.ProjectID, f.AgentID) + if err != nil { + return TriggeredCompactionHandoff{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.BeginCompaction( + ctx, + agentexecution.BeginCompactionInput{ + Failure: executionFailure(f), + SourceEnd: input.SourceEventSequenceEnd, + }, + ) + if err != nil { + return TriggeredCompactionHandoff{}, err + } + if err := applyAdmissionDestination(ctx, + unit, + input.Failure.ProjectID, + input.Failure.AgentID, + result.Admission); err != nil { + return TriggeredCompactionHandoff{}, err + } + parent, err := loadModelCallContextByID( + ctx, + dbsqlc.New(unit.DB()), + f.ProjectID, + f.AgentID, + input.ParentContextID, + ) + if err != nil { + return TriggeredCompactionHandoff{}, err + } + claim, err := executionClaim(ctx, unit, f.ProjectID, f.AgentID, result.Model) + if err != nil { + return TriggeredCompactionHandoff{}, err + } + if err = unit.Commit(ctx, "begin compaction"); err != nil { + return TriggeredCompactionHandoff{}, err + } + return TriggeredCompactionHandoff{ + ParentContext: parent, + CompactionCall: claim, + BoundaryPreempted: result.Preempted, + }, nil + }) } func (s *Store) ReplaceCompactionSource( ctx context.Context, input ReplaceCompactionSourceInput, ) (ReplaceCompactionSourceResult, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return ReplaceCompactionSourceResult{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - failure := agentexecution.ModelFailure{ - ContextID: input.ModelCallContextID, - RuntimeLockID: input.RuntimeLockID, - ErrorKind: input.ErrorKind, - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: input.ErrorDetails, - Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, APIVariant: input.APIVariant, - RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, Usage: input.Usage, - Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}, - } - result, err := h.BeginCompaction( - ctx, - agentexecution.BeginCompactionInput{Failure: failure, SourceEnd: input.NextSourceEventSequenceEnd}, - ) - if err != nil { - return ReplaceCompactionSourceResult{}, err - } - if err := applyAdmissionDestination(ctx, - unit, - input.ProjectID, - input.AgentID, - result.Admission); err != nil { - return ReplaceCompactionSourceResult{}, err - } - claim, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, result.Model) - if err != nil { - return ReplaceCompactionSourceResult{}, err - } - if err = unit.Commit(ctx, "reduce compaction source"); err != nil { - return ReplaceCompactionSourceResult{}, err - } - return ReplaceCompactionSourceResult{CompactionCall: claim, BoundaryPreempted: result.Preempted}, nil + return agentexecution.Retry(ctx, + "ReplaceCompactionSource", + func(ctx context.Context) (ReplaceCompactionSourceResult, + error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return ReplaceCompactionSourceResult{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + failure := agentexecution.ModelFailure{ + ContextID: input.ModelCallContextID, + RuntimeLockID: input.RuntimeLockID, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, APIVariant: input.APIVariant, + RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, Usage: input.Usage, + Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}, + } + result, err := h.BeginCompaction( + ctx, + agentexecution.BeginCompactionInput{Failure: failure, SourceEnd: input.NextSourceEventSequenceEnd}, + ) + if err != nil { + return ReplaceCompactionSourceResult{}, err + } + if err := applyAdmissionDestination(ctx, + unit, + input.ProjectID, + input.AgentID, + result.Admission); err != nil { + return ReplaceCompactionSourceResult{}, err + } + claim, err := executionClaim(ctx, unit, input.ProjectID, input.AgentID, result.Model) + if err != nil { + return ReplaceCompactionSourceResult{}, err + } + if err = unit.Commit(ctx, "reduce compaction source"); err != nil { + return ReplaceCompactionSourceResult{}, err + } + return ReplaceCompactionSourceResult{CompactionCall: claim, BoundaryPreempted: result.Preempted}, nil + }) } diff --git a/internal/storage/executionstore/model_contexts_lifecycle.go b/internal/storage/executionstore/model_contexts_lifecycle.go index 51ff023c1..c386ebe01 100644 --- a/internal/storage/executionstore/model_contexts_lifecycle.go +++ b/internal/storage/executionstore/model_contexts_lifecycle.go @@ -4,6 +4,7 @@ import ( "context" "errors" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" ) @@ -11,32 +12,37 @@ func (s *Store) RecordRetryableModelCallFailure( ctx context.Context, input RecordRecoverableModelCallFailureInput, ) (ModelCallContextRecord, error) { - if input.RecoveryKind == "" { - input.RecoveryKind = ModelCallRecoveryRetry - } - if input.RecoveryKind != ModelCallRecoveryRetry { - return ModelCallContextRecord{}, errors.New("retry recovery is required") - } - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return ModelCallContextRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - if _, err = h.FailModel(ctx, executionFailure(input)); err != nil { - return ModelCallContextRecord{}, err - } - record, err := loadModelCallContextByID( - ctx, - dbsqlc.New(unit.DB()), - input.ProjectID, - input.AgentID, - input.ModelCallContextID, - ) - if err != nil { - return ModelCallContextRecord{}, err - } - if err = unit.Commit(ctx, "record retryable failure"); err != nil { - return ModelCallContextRecord{}, err - } - return record, nil + return agentexecution.Retry(ctx, + "RecordRetryableModelCallFailure", + func(ctx context.Context) (ModelCallContextRecord, + error) { + if input.RecoveryKind == "" { + input.RecoveryKind = ModelCallRecoveryRetry + } + if input.RecoveryKind != ModelCallRecoveryRetry { + return ModelCallContextRecord{}, errors.New("retry recovery is required") + } + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return ModelCallContextRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if _, err = h.FailModel(ctx, executionFailure(input)); err != nil { + return ModelCallContextRecord{}, err + } + record, err := loadModelCallContextByID( + ctx, + dbsqlc.New(unit.DB()), + input.ProjectID, + input.AgentID, + input.ModelCallContextID, + ) + if err != nil { + return ModelCallContextRecord{}, err + } + if err = unit.Commit(ctx, "record retryable failure"); err != nil { + return ModelCallContextRecord{}, err + } + return record, nil + }) } diff --git a/internal/storage/executionstore/model_contexts_preparation.go b/internal/storage/executionstore/model_contexts_preparation.go index 4d2fb906c..487bfef91 100644 --- a/internal/storage/executionstore/model_contexts_preparation.go +++ b/internal/storage/executionstore/model_contexts_preparation.go @@ -5,6 +5,7 @@ import ( "errors" "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) type PrepareNormalModelCallInput struct { @@ -26,20 +27,24 @@ func (s *Store) PrepareNormalModelCall( ctx context.Context, input PrepareNormalModelCallInput, ) (PreparedNormalModelCall, error) { - if input.RuntimeLockID == uuid.Nil || len(input.OpeningInputIDs) == 0 { - return PreparedNormalModelCall{}, errors.New("runtime and opening inputs are required") - } - unit, _, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return PreparedNormalModelCall{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - result, err := prepareExecutionModel(ctx, unit, input) - if err != nil { - return PreparedNormalModelCall{}, err - } - if err = unit.Commit(ctx, "prepare model call"); err != nil { - return PreparedNormalModelCall{}, err - } - return result, nil + return agentexecution.Retry(ctx, + "PrepareNormalModelCall", + func(ctx context.Context) (PreparedNormalModelCall, error) { + if input.RuntimeLockID == uuid.Nil || len(input.OpeningInputIDs) == 0 { + return PreparedNormalModelCall{}, errors.New("runtime and opening inputs are required") + } + unit, _, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return PreparedNormalModelCall{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := prepareExecutionModel(ctx, unit, input) + if err != nil { + return PreparedNormalModelCall{}, err + } + if err = unit.Commit(ctx, "prepare model call"); err != nil { + return PreparedNormalModelCall{}, err + } + return result, nil + }) } diff --git a/internal/storage/executionstore/model_contexts_preparation_integration_test.go b/internal/storage/executionstore/model_contexts_preparation_integration_test.go index c99b816e2..358950498 100644 --- a/internal/storage/executionstore/model_contexts_preparation_integration_test.go +++ b/internal/storage/executionstore/model_contexts_preparation_integration_test.go @@ -7,7 +7,6 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/executionstore" - "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/storeerr" "github.com/omnara-ai/omnara/internal/testutil/integrationdb" @@ -24,7 +23,7 @@ func TestPrepareNormalModelCallUsesSnapshotAfterLockWait(t *testing.T) { _, err = blocker.LockAgent( ctx, dbsqlc.LockAgentInProjectParams{ProjectID: testProjectID, ID: f.AgentID}, - agentexecution.ExternalAuthority{}, + nil, ) require.NoError(t, err) var pid int32 diff --git a/internal/storage/executionstore/model_contexts_replay_integration_test.go b/internal/storage/executionstore/model_contexts_replay_integration_test.go index 7486cef2a..f03e2c0c9 100644 --- a/internal/storage/executionstore/model_contexts_replay_integration_test.go +++ b/internal/storage/executionstore/model_contexts_replay_integration_test.go @@ -42,9 +42,12 @@ INSERT INTO model_call_contexts( ) VALUES ($1, $2, $3, 'normal', 1, $4, $5, $6, $7, 'started', statement_timestamp(), (SELECT current_turn_id FROM agent_execution_state WHERE agent_id=$3), (SELECT -array_agg(agent_input_id ORDER BY sequence) FROM agent_events WHERE agent_id=$3 AND -is_opening_event AND event_kind='agent_input'), (SELECT min(sequence) FROM agent_events -WHERE agent_id=$3 AND is_opening_event)) +array_agg(i.id ORDER BY e.sequence) FROM agent_events e +JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id +WHERE e.agent_id=$3 AND e.is_opening_event AND i.input_kind='content'), +(SELECT min(e.sequence) FROM agent_events e +JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id +WHERE e.agent_id=$3 AND e.is_opening_event AND i.input_kind='content')) RETURNING id `, testOrgID, testProjectID, fixture.AgentID, agent.CurrentConfigID, revisionID, frontier, fixture.Lock.ID).Scan(&contextID); err != nil { @@ -63,9 +66,12 @@ INSERT INTO model_call_contexts( ) VALUES ($1, $2, $3, 'compaction', 1, $4, $5, $6, 1, $7, 'started', statement_timestamp(), (SELECT current_turn_id FROM agent_execution_state WHERE agent_id=$3), (SELECT -array_agg(agent_input_id ORDER BY sequence) FROM agent_events WHERE agent_id=$3 AND -is_opening_event AND event_kind='agent_input'), (SELECT min(sequence) FROM agent_events -WHERE agent_id=$3 AND is_opening_event)) +array_agg(i.id ORDER BY e.sequence) FROM agent_events e +JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id +WHERE e.agent_id=$3 AND e.is_opening_event AND i.input_kind='content'), +(SELECT min(e.sequence) FROM agent_events e +JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id +WHERE e.agent_id=$3 AND e.is_opening_event AND i.input_kind='content')) RETURNING id `, testOrgID, testProjectID, fixture.AgentID, agent.CurrentConfigID, revisionID, frontier, fixture.Lock.ID).Scan(&contextID); err != nil { diff --git a/internal/storage/executionstore/model_outputs_terminal_error.go b/internal/storage/executionstore/model_outputs_terminal_error.go index bbb165344..3e661e651 100644 --- a/internal/storage/executionstore/model_outputs_terminal_error.go +++ b/internal/storage/executionstore/model_outputs_terminal_error.go @@ -35,46 +35,51 @@ func (s *Store) RecordModelCallErrorAndCompleteContext( ctx context.Context, input RecordModelCallErrorAndCompleteContextInput, ) (events.Event, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return events.Event{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - _, err = h.FailModel( - ctx, - agentexecution.ModelFailure{ContextID: input.ModelCallContextID, RuntimeLockID: input.RuntimeLockID, - ErrorKind: input.ErrorKind, - ErrorCode: input.ErrorCode, - ErrorMessage: input.ErrorMessage, - ErrorDetails: input.ErrorDetails, - ServedModel: input.ServedProviderModelSlug, - Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, - APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, - Usage: input.Usage, Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}}, - ) - if err != nil { - return events.Event{}, err - } - q := dbsqlc.New(unit.DB()) - output, err := q.GetModelOutputByModelContext( - ctx, - dbsqlc.GetModelOutputByModelContextParams{ProjectID: input.ProjectID, - AgentID: input.AgentID, ModelCallContextID: input.ModelCallContextID}, - ) - if err != nil { - return events.Event{}, err - } - row, err := q.GetTypedAgentEventByModelOutput(ctx, dbsqlc.GetTypedAgentEventByModelOutputParams{ - ProjectID: input.ProjectID, AgentID: input.AgentID, ModelOutputID: &output.ID}) - if err != nil { - return events.Event{}, err - } - event, err := typedAgentEventFromModelOutputSQLC(row) - if err != nil { - return events.Event{}, err - } - if err = unit.Commit(ctx, "record terminal model failure"); err != nil { - return events.Event{}, err - } - return event.Event, nil + return agentexecution.Retry(ctx, + "RecordModelCallErrorAndCompleteContext", + func(ctx context.Context) (events.Event, + error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return events.Event{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + _, err = h.FailModel( + ctx, + agentexecution.ModelFailure{ContextID: input.ModelCallContextID, RuntimeLockID: input.RuntimeLockID, + ErrorKind: input.ErrorKind, + ErrorCode: input.ErrorCode, + ErrorMessage: input.ErrorMessage, + ErrorDetails: input.ErrorDetails, + ServedModel: input.ServedProviderModelSlug, + Evidence: agentexecution.ModelEvidence{APIFormat: input.APIFormat, + APIVariant: input.APIVariant, RequestID: input.ProviderRequestID, ResponseID: input.ProviderResponseID, + Usage: input.Usage, Cost: input.ProviderReportedCostUSD, Metadata: input.ProviderMetadata}}, + ) + if err != nil { + return events.Event{}, err + } + q := dbsqlc.New(unit.DB()) + output, err := q.GetModelOutputByModelContext( + ctx, + dbsqlc.GetModelOutputByModelContextParams{ProjectID: input.ProjectID, + AgentID: input.AgentID, ModelCallContextID: input.ModelCallContextID}, + ) + if err != nil { + return events.Event{}, err + } + row, err := q.GetTypedAgentEventByModelOutput(ctx, dbsqlc.GetTypedAgentEventByModelOutputParams{ + ProjectID: input.ProjectID, AgentID: input.AgentID, ModelOutputID: &output.ID}) + if err != nil { + return events.Event{}, err + } + event, err := typedAgentEventFromModelOutputSQLC(row) + if err != nil { + return events.Event{}, err + } + if err = unit.Commit(ctx, "record terminal model failure"); err != nil { + return events.Event{}, err + } + return event.Event, nil + }) } diff --git a/internal/storage/executionstore/model_outputs_workflow.go b/internal/storage/executionstore/model_outputs_workflow.go index d486dc771..6c83153a4 100644 --- a/internal/storage/executionstore/model_outputs_workflow.go +++ b/internal/storage/executionstore/model_outputs_workflow.go @@ -59,47 +59,63 @@ func (s *Store) RecordModelOutputAndCompleteContext( ctx context.Context, input RecordModelOutputAndCompleteContextInput, ) (events.Event, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return events.Event{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - result, err := h.AcceptOutput(ctx, executionOutput(input, nil)) - if err != nil { - return events.Event{}, err - } - if err = unit.Commit(ctx, "accept model output"); err != nil { - return events.Event{}, err - } - return executionEvent(input.AgentID, result.Event, events.KindModelOutput), nil + return agentexecution.Retry(ctx, + "RecordModelOutputAndCompleteContext", + func(ctx context.Context) (events.Event, + error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return events.Event{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.AcceptOutput(ctx, executionOutput(input, nil)) + if err != nil { + return events.Event{}, err + } + if err = unit.Commit(ctx, "accept model output"); err != nil { + return events.Event{}, err + } + return executionEvent(input.AgentID, result.Event, events.KindModelOutput), nil + }) } func (s *Store) RecordToolCallSourceAndCompleteContext( ctx context.Context, input RecordToolCallSourceAndCompleteContextInput, ) (events.Event, []ToolCallRecord, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return events.Event{}, nil, err - } - defer func() { _ = unit.Rollback(ctx) }() - result, err := h.AcceptOutput(ctx, executionOutput(RecordModelOutputAndCompleteContextInput{ - ProjectID: input.ProjectID, AgentID: input.AgentID, RuntimeLockID: input.RuntimeLockID, - ModelCallContextID: input.ModelCallContextID, ProviderRequestID: input.ProviderRequestID, - ProviderResponse: input.ProviderResponse}, input.ToolCallBindings)) - if err != nil { - return events.Event{}, nil, err - } - records := make([]ToolCallRecord, 0, len(result.Tools)) - for _, call := range result.Tools { - record, loadErr := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, call.ID) - if loadErr != nil { - return events.Event{}, nil, loadErr - } - records = append(records, record) - } - if err = unit.Commit(ctx, "accept model tool output"); err != nil { - return events.Event{}, nil, err + type executionResult struct { + event events.Event + tools []ToolCallRecord } - return executionEvent(input.AgentID, result.Event, events.KindModelOutput), records, nil + result, + err := agentexecution.Retry(ctx, + "RecordToolCallSourceAndCompleteContext", + func(ctx context.Context) (executionResult, + error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return executionResult{events.Event{}, nil}, err + } + defer func() { _ = unit.Rollback(ctx) }() + result, err := h.AcceptOutput(ctx, executionOutput(RecordModelOutputAndCompleteContextInput{ + ProjectID: input.ProjectID, AgentID: input.AgentID, RuntimeLockID: input.RuntimeLockID, + ModelCallContextID: input.ModelCallContextID, ProviderRequestID: input.ProviderRequestID, + ProviderResponse: input.ProviderResponse}, input.ToolCallBindings)) + if err != nil { + return executionResult{events.Event{}, nil}, err + } + records := make([]ToolCallRecord, 0, len(result.Tools)) + for _, call := range result.Tools { + record, loadErr := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, call.ID) + if loadErr != nil { + return executionResult{events.Event{}, nil}, loadErr + } + records = append(records, record) + } + if err = unit.Commit(ctx, "accept model tool output"); err != nil { + return executionResult{events.Event{}, nil}, err + } + return executionResult{executionEvent(input.AgentID, result.Event, events.KindModelOutput), records}, nil + }) + return result.event, result.tools, err } diff --git a/internal/storage/executionstore/model_work_transition.go b/internal/storage/executionstore/model_work_transition.go index bed2d5d62..8531d5e00 100644 --- a/internal/storage/executionstore/model_work_transition.go +++ b/internal/storage/executionstore/model_work_transition.go @@ -2,6 +2,9 @@ package executionstore import ( "context" + "errors" + + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) type AcceptModelOutputAndAdvanceInput struct { @@ -14,44 +17,69 @@ type AcceptModelOutputAndAdvanceInput struct { type OwnedAgentWorkTransition struct { Work ClaimedAgentWork Continued bool + Fallback error } func (s *Store) AcceptModelOutputAndAdvance( ctx context.Context, input AcceptModelOutputAndAdvanceInput, ) (OwnedAgentWorkTransition, error) { + return agentexecution.Retry(ctx, + "accept model output", + func(ctx context.Context) (OwnedAgentWorkTransition, error) { + return s.acceptModelOutput(ctx, input, true) + }) +} + +func (s *Store) acceptModelOutput(ctx context.Context, + input AcceptModelOutputAndAdvanceInput, + advance bool) (OwnedAgentWorkTransition, + error) { output := input.Output unit, h, err := beginExecution(ctx, s, output.ProjectID, output.AgentID) if err != nil { return OwnedAgentWorkTransition{}, err } defer func() { _ = unit.Rollback(ctx) }() - accepted, work, err := h.AcceptAndAdvance( - ctx, - executionOutput(output, input.ToolCallBindings), - input.AllowModelWork, - input.PrepareModel, - ) + accepted, err := h.AcceptOutput(ctx, executionOutput(output, input.ToolCallBindings)) if err != nil { return OwnedAgentWorkTransition{}, err } result := OwnedAgentWorkTransition{} - if accepted.Created { - if err := applyAdmissionDestination(ctx, - unit, - input.Output.ProjectID, - input.Output.AgentID, - work.Admission); err != nil { - return result, err - } - result.Work, err = shapeOwnedWork(ctx, unit, output.ProjectID, output.AgentID, work) - if err != nil { + if advance && accepted.Created { + err = unit.Savepoint(ctx, func(_ *agentexecution.Unit) error { + work, err := h.Advance(ctx, output.RuntimeLockID, input.AllowModelWork, input.PrepareModel) + if err != nil { + return err + } + if err = applyAdmissionDestination(ctx, unit, output.ProjectID, output.AgentID, work.Admission); err != nil { + return err + } + shaped, err := shapeOwnedWork(ctx, unit, output.ProjectID, output.AgentID, work) + if err != nil { + return err + } + result = OwnedAgentWorkTransition{Work: shaped, Continued: !work.Released} + return nil + }) + } + if errors.Is(err, agentexecution.ErrLockPlan) { + return OwnedAgentWorkTransition{}, err + } + if !advance || !accepted.Created || err != nil { + result = OwnedAgentWorkTransition{Fallback: err} + if err = h.Release(ctx, output.RuntimeLockID); err != nil { return result, err } - result.Continued = !work.Released } if err = unit.Commit(ctx, "accept model output and advance"); err != nil { - return OwnedAgentWorkTransition{}, err + if !advance || errors.Is(err, agentexecution.ErrLockPlan) { + return result, err + } + _ = unit.Rollback(ctx) + fallback, acceptErr := s.acceptModelOutput(ctx, input, false) + fallback.Fallback = errors.Join(result.Fallback, err, fallback.Fallback) + return fallback, acceptErr } return result, nil } diff --git a/internal/storage/executionstore/model_work_transition_integration_test.go b/internal/storage/executionstore/model_work_transition_integration_test.go index 672207ac9..15eb807c8 100644 --- a/internal/storage/executionstore/model_work_transition_integration_test.go +++ b/internal/storage/executionstore/model_work_transition_integration_test.go @@ -7,10 +7,12 @@ import ( "testing" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" "github.com/omnara-ai/omnara/internal/modelenvelope" "github.com/omnara-ai/omnara/internal/modelprotocol" "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" "github.com/stretchr/testify/require" ) @@ -145,31 +147,10 @@ func TestAcceptModelOutputAndAdvanceFencing(t *testing.T) { } } -func TestAcceptModelOutputAndAdvanceRollbackAndCrashRecovery(t *testing.T) { +func TestAcceptModelOutputAndAdvanceCrashRecovery(t *testing.T) { ctx := t.Context() f, input := fusedModelFixture(t, ctx) input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonMaxTokens - _, err := f.Store.pool.Exec(ctx, `CREATE FUNCTION reject_fused_preparation() RETURNS trigger LANGUAGE plpgsql AS $$ - BEGIN RAISE EXCEPTION 'reject fused preparation'; END $$; - CREATE CONSTRAINT TRIGGER reject_fused_preparation AFTER INSERT ON model_call_contexts - DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION reject_fused_preparation();`) - require.NoError(t, err) - _, err = f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) - require.ErrorContains(t, err, "commit accept model output and advance") - _, found, err := f.Store.Execution().GetModelOutputForContext( - ctx, testProjectID, f.AgentID, input.Output.ModelCallContextID, - ) - require.NoError(t, err) - require.False(t, found) - original, found, err := f.Store.Execution().GetModelCallContext( - ctx, testProjectID, f.AgentID, input.Output.ModelCallContextID, - ) - require.NoError(t, err) - require.True(t, found) - require.Equal(t, executionstore.ModelCallContextStarted, original.State) - requireOwnedLease(t, ctx, f, true) - _, err = f.Store.pool.Exec(ctx, `DROP TRIGGER reject_fused_preparation ON model_call_contexts`) - require.NoError(t, err) transition, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) require.NoError(t, err) require.True(t, transition.Continued) @@ -212,3 +193,71 @@ func TestAcceptModelOutputAndAdvanceReplayRecoversPreparedWorkWithoutDuplicating ).Scan(&outputs)) require.Equal(t, 1, outputs) } + +func TestAcceptModelOutputFallbackRechecksFence(t *testing.T) { + for _, mode := range []string{"cancel", "expired"} { + t.Run(mode, func(t *testing.T) { + ctx := t.Context() + f, input := fusedModelFixture(t, ctx) + input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonMaxTokens + gate, err := f.Store.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = gate.Rollback(ctx) }() + var gatePID int32 + require.NoError(t, gate.QueryRow(ctx, `SELECT pg_backend_pid()`).Scan(&gatePID)) + _, err = gate.Exec(ctx, `SELECT pg_advisory_xact_lock(63109)`) + require.NoError(t, err) + _, err = f.Store.pool.Exec(ctx, ` +CREATE FUNCTION reject_successor() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF EXISTS(SELECT 1 FROM model_outputs WHERE agent_id=NEW.agent_id) THEN + PERFORM pg_advisory_xact_lock(63109); + RAISE EXCEPTION 'reject successor'; + END IF; + RETURN NEW; +END $$; +CREATE CONSTRAINT TRIGGER reject_successor AFTER INSERT ON model_call_contexts + DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION reject_successor()`) + require.NoError(t, err) + type acceptance struct { + transition executionstore.OwnedAgentWorkTransition + err error + } + accepted := make(chan acceptance, 1) + go func() { + transition, err := f.Store.Execution().AcceptModelOutputAndAdvance(ctx, input) + accepted <- acceptance{transition, err} + }() + integrationdb.WaitForLockWaitBlockedBy(t, ctx, f.Store.pool, "commit", gatePID) + var acceptPID int32 + require.NoError(t, f.Store.pool.QueryRow(ctx, `SELECT pid FROM pg_stat_activity +WHERE datname=current_database() AND $1=ANY(pg_blocking_pids(pid))`, gatePID).Scan(&acceptPID)) + fenced := make(chan error, 1) + actor := mustOmnaraActorParams(t, f.UserID) + go func() { + if mode == "cancel" { + _, err := f.Store.Execution().CancelAgent(ctx, executionstore.CancelAgentInput{ + ProjectID: testProjectID, AgentID: f.AgentID, Actor: actor, + }) + fenced <- err + } else { + _, err := f.Store.pool.Exec(ctx, `UPDATE agent_runtime_locks +SET renewed_at=started_at,lease_expires_at=statement_timestamp()-interval '1 microsecond' WHERE id=$1`, f.Lock.ID) + fenced <- err + } + }() + integrationdb.WaitForLockWaitBlockedBy(t, ctx, f.Store.pool, "", acceptPID) + require.NoError(t, gate.Rollback(ctx)) + require.NoError(t, <-fenced) + result := <-accepted + require.ErrorIs(t, result.err, storeerr.ErrRuntimeLockInactive) + var cause *pgconn.PgError + require.ErrorAs(t, result.transition.Fallback, &cause) + require.Equal(t, "reject successor", cause.Message) + var outputs int + require.NoError(t, f.Store.pool.QueryRow(ctx, + `SELECT count(*) FROM model_outputs WHERE agent_id=$1`, f.AgentID).Scan(&outputs)) + require.Zero(t, outputs) + }) + } +} diff --git a/internal/storage/executionstore/process_action_completion_store.go b/internal/storage/executionstore/process_action_completion_store.go index 6906ca58d..dfdcc5bb5 100644 --- a/internal/storage/executionstore/process_action_completion_store.go +++ b/internal/storage/executionstore/process_action_completion_store.go @@ -229,34 +229,38 @@ func (s *Store) completeDaemonProcessAction( "project, agent, process, and action are required", ) } - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonProcessActionReportApplication{}, fmt.Errorf("begin complete process action: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + return agentexecution.Retry(ctx, + "completeDaemonProcessAction", + func(ctx context.Context) (DaemonProcessActionReportApplication, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return DaemonProcessActionReportApplication{}, fmt.Errorf("begin complete process action: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() - qtx := dbsqlc.New(tx) - if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { - return DaemonProcessActionReportApplication{}, err - } - completion := processActionCompletionInput{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ProcessID: input.ProcessID, - ID: input.ID, - StateReasonCode: input.StateReasonCode, - StateReasonMessage: input.StateReasonMessage, - Result: input.Result, - } - application, err := completeDaemonProcessActionTx(ctx, unit, qtx, completion, state) - if err != nil { - return DaemonProcessActionReportApplication{}, err - } - if err := unit.Commit(ctx, "complete process action"); err != nil { - return DaemonProcessActionReportApplication{}, err - } - return application, nil + qtx := dbsqlc.New(tx) + if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { + return DaemonProcessActionReportApplication{}, err + } + completion := processActionCompletionInput{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + ProcessID: input.ProcessID, + ID: input.ID, + StateReasonCode: input.StateReasonCode, + StateReasonMessage: input.StateReasonMessage, + Result: input.Result, + } + application, err := completeDaemonProcessActionTx(ctx, unit, qtx, completion, state) + if err != nil { + return DaemonProcessActionReportApplication{}, err + } + if err := unit.Commit(ctx, "complete process action"); err != nil { + return DaemonProcessActionReportApplication{}, err + } + return application, nil + }) } func completeDaemonProcessActionTx( @@ -269,7 +273,7 @@ func completeDaemonProcessActionTx( _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ ProjectID: input.ProjectID, ID: input.AgentID, - }, agentexecution.ExternalAuthority{}) + }, nil) if err != nil { return DaemonProcessActionReportApplication{}, fmt.Errorf( "lock agent for process action completion: %w", diff --git a/internal/storage/executionstore/process_cancellation_archive_integration_test.go b/internal/storage/executionstore/process_cancellation_archive_integration_test.go index b7f9e0443..4b9ffbe02 100644 --- a/internal/storage/executionstore/process_cancellation_archive_integration_test.go +++ b/internal/storage/executionstore/process_cancellation_archive_integration_test.go @@ -492,9 +492,9 @@ func TestArchiveAgentAtomicallyStopsDurableModelCallWork(t *testing.T) { contextRecord, ) } - if !test.recordRetry && contextRecord.CompletedAt.Before(*archived.ArchivedAt) { + if !test.recordRetry && contextRecord.CompletedAt.After(*archived.ArchivedAt) { t.Fatalf( - "archive-triggered model cancellation predates the agent archive: agent=%+v context=%+v", + "archive-triggered model cancellation follows the agent archive: agent=%+v context=%+v", archived, contextRecord, ) diff --git a/internal/storage/executionstore/process_daemon_completion_store.go b/internal/storage/executionstore/process_daemon_completion_store.go index f1f761840..1ac5a6dcc 100644 --- a/internal/storage/executionstore/process_daemon_completion_store.go +++ b/internal/storage/executionstore/process_daemon_completion_store.go @@ -33,88 +33,92 @@ func (s *Store) MarkProcessStarted( ) } input.SourceStartedAt = canonicalSourceTime(input.SourceStartedAt) - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("begin mark process started: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - qtx := dbsqlc.New(tx) - if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { - return DaemonProcessReportApplication{}, err - } - _, err = unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, - ID: input.AgentID, - }, agentexecution.ExternalAuthority{}) - if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf( - "lock agent for process start observation: %w", - err, - ) - } - row, err := qtx.MarkProcessStarted( - ctx, - dbsqlc.MarkProcessStartedParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - MachineID: input.Authority.MachineID, - SourceStartedAt: input.SourceStartedAt, - }, - ) - var record ProcessRecord - if errors.Is(err, pgx.ErrNoRows) { - record, err = daemonProcessForReportTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.Authority, - input.ID, - false, - ) - if err != nil { - return DaemonProcessReportApplication{}, err - } - if !isProcessTerminal(record.State) { - return DaemonProcessReportApplication{}, storeerr.ErrDaemonRuntimeUnregistered - } - } else if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("mark process started: %w", err) - } else { - record = processRecordFromStartedSQLC(row) - } - resultCommitted := false - if record.ToolCallID != uuid.Nil && record.ExecutionSpec.FileTransfer == nil { - h, err := unit.Handle(input.ProjectID, input.AgentID) - if err != nil { - return DaemonProcessReportApplication{}, err - } - completed, err := h.CompleteProcess( - ctx, - agentexecution.ProcessResult{ID: record.ID, Started: true, Observed: input.Result}, - ) - if err != nil { - return DaemonProcessReportApplication{}, err - } - resultCommitted = completed.Matches - row, err := qtx.GetProcess( - ctx, - dbsqlc.GetProcessParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: record.ID}, - ) - if err != nil { - return DaemonProcessReportApplication{}, err - } - record = processRecordFromSQLC(row) - } - if err := unit.Commit(ctx, "mark process started"); err != nil { - return DaemonProcessReportApplication{}, err - } - return DaemonProcessReportApplication{ - Process: record, - ToolResultCommitted: resultCommitted, - }, nil + return agentexecution.Retry(ctx, + "MarkProcessStarted", + func(ctx context.Context) (DaemonProcessReportApplication, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return DaemonProcessReportApplication{}, fmt.Errorf("begin mark process started: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) + if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { + return DaemonProcessReportApplication{}, err + } + _, err = unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: input.ProjectID, + ID: input.AgentID, + }, nil) + if err != nil { + return DaemonProcessReportApplication{}, fmt.Errorf( + "lock agent for process start observation: %w", + err, + ) + } + row, err := qtx.MarkProcessStarted( + ctx, + dbsqlc.MarkProcessStartedParams{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + ID: input.ID, + MachineID: input.Authority.MachineID, + SourceStartedAt: input.SourceStartedAt, + }, + ) + var record ProcessRecord + if errors.Is(err, pgx.ErrNoRows) { + record, err = daemonProcessForReportTx( + ctx, + qtx, + input.ProjectID, + input.AgentID, + input.Authority, + input.ID, + false, + ) + if err != nil { + return DaemonProcessReportApplication{}, err + } + if !isProcessTerminal(record.State) { + return DaemonProcessReportApplication{}, storeerr.ErrDaemonRuntimeUnregistered + } + } else if err != nil { + return DaemonProcessReportApplication{}, fmt.Errorf("mark process started: %w", err) + } else { + record = processRecordFromStartedSQLC(row) + } + resultCommitted := false + if record.ToolCallID != uuid.Nil && record.ExecutionSpec.FileTransfer == nil { + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return DaemonProcessReportApplication{}, err + } + completed, err := h.CompleteProcess( + ctx, + agentexecution.ProcessResult{ID: record.ID, Started: true, Observed: input.Result}, + ) + if err != nil { + return DaemonProcessReportApplication{}, err + } + resultCommitted = completed.Matches + row, err := qtx.GetProcess( + ctx, + dbsqlc.GetProcessParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: record.ID}, + ) + if err != nil { + return DaemonProcessReportApplication{}, err + } + record = processRecordFromSQLC(row) + } + if err := unit.Commit(ctx, "mark process started"); err != nil { + return DaemonProcessReportApplication{}, err + } + return DaemonProcessReportApplication{ + Process: record, + ToolResultCommitted: resultCommitted, + }, nil + }) } func (s *Store) CompleteDaemonProcess( @@ -147,167 +151,171 @@ func (s *Store) CompleteDaemonProcess( if !input.SourceEndedAt.IsZero() { input.SourceEndedAt = canonicalSourceTime(input.SourceEndedAt) } - unit, err := s.cell.Begin(ctx) - if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("begin complete daemon process: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - - qtx := dbsqlc.New(tx) - if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { - return DaemonProcessReportApplication{}, err - } - _, err = unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: input.ProjectID, - ID: input.AgentID, - }, agentexecution.ExternalAuthority{}) - if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf( - "lock agent for daemon process completion: %w", - err, - ) - } - var startedAt *time.Time - if !input.SourceStartedAt.IsZero() { - startedAt = &input.SourceStartedAt - } - var endedAt *time.Time - if !input.SourceEndedAt.IsZero() { - endedAt = &input.SourceEndedAt - } - row, err := qtx.CompleteDaemonObservedProcess( - ctx, - dbsqlc.CompleteDaemonObservedProcessParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - MachineID: input.Authority.MachineID, - State: string(input.State), - SourceStartedAt: startedAt, - SourceEndedAt: endedAt, - ExitCode: storeutil.Int32Ptr(input.ExitCode), - ExitSignal: input.ExitSignal, - StateReasonCode: storeutil.TextFromEmpty(input.StateReasonCode), - StateReasonMessage: input.StateReasonMessage, - StorageExhausted: input.StorageExhausted, - }, - ) - processUpdated := err == nil - var record ProcessRecord - if errors.Is(err, pgx.ErrNoRows) { - record, err = daemonProcessForReportTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - input.Authority, - input.ID, - input.StorageExhausted, - ) - if err != nil { - return DaemonProcessReportApplication{}, err - } - if !isProcessTerminal(record.State) { - return DaemonProcessReportApplication{}, storeerr.ErrDaemonRuntimeUnregistered - } - } else if err != nil { - return DaemonProcessReportApplication{}, fmt.Errorf("complete daemon process: %w", err) - } else { - record = processRecordFromCompleteSQLC(row) - } - reportMatchesProcess := processUpdated || daemonTerminalReportMatchesRecord(record, input) - resultCommitted := false - if reportMatchesProcess && record.ToolCallID != uuid.Nil { - request := agentexecution.ProcessResult{ID: record.ID, Observed: input.Result} - if record.ExecutionSpec.FileTransfer != nil { - _, result, err := processresult.Terminal(record.resultFacts()) + return agentexecution.Retry(ctx, + "CompleteDaemonProcess", + func(ctx context.Context) (DaemonProcessReportApplication, error) { + unit, err := s.cell.Begin(ctx) if err != nil { + return DaemonProcessReportApplication{}, fmt.Errorf("begin complete daemon process: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + + qtx := dbsqlc.New(tx) + if err := requireReportableDaemonRuntimeAuthorityTx(ctx, qtx, input.Authority); err != nil { return DaemonProcessReportApplication{}, err } - if len(input.Result) > 0 && string(input.Result) != "null" { - result, err = commandTerminalToolResult(record.ID, input.Result) + _, err = unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: input.ProjectID, + ID: input.AgentID, + }, nil) + if err != nil { + return DaemonProcessReportApplication{}, fmt.Errorf( + "lock agent for daemon process completion: %w", + err, + ) + } + var startedAt *time.Time + if !input.SourceStartedAt.IsZero() { + startedAt = &input.SourceStartedAt + } + var endedAt *time.Time + if !input.SourceEndedAt.IsZero() { + endedAt = &input.SourceEndedAt + } + row, err := qtx.CompleteDaemonObservedProcess( + ctx, + dbsqlc.CompleteDaemonObservedProcessParams{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + ID: input.ID, + MachineID: input.Authority.MachineID, + State: string(input.State), + SourceStartedAt: startedAt, + SourceEndedAt: endedAt, + ExitCode: storeutil.Int32Ptr(input.ExitCode), + ExitSignal: input.ExitSignal, + StateReasonCode: storeutil.TextFromEmpty(input.StateReasonCode), + StateReasonMessage: input.StateReasonMessage, + StorageExhausted: input.StorageExhausted, + }, + ) + processUpdated := err == nil + var record ProcessRecord + if errors.Is(err, pgx.ErrNoRows) { + record, err = daemonProcessForReportTx( + ctx, + qtx, + input.ProjectID, + input.AgentID, + input.Authority, + input.ID, + input.StorageExhausted, + ) if err != nil { return DaemonProcessReportApplication{}, err } + if !isProcessTerminal(record.State) { + return DaemonProcessReportApplication{}, storeerr.ErrDaemonRuntimeUnregistered + } + } else if err != nil { + return DaemonProcessReportApplication{}, fmt.Errorf("complete daemon process: %w", err) + } else { + record = processRecordFromCompleteSQLC(row) } - outcome, raw, err := fileTransferToolResultContentParts(ctx, qtx, record, result) - if err != nil { - return DaemonProcessReportApplication{}, err + reportMatchesProcess := processUpdated || daemonTerminalReportMatchesRecord(record, input) + resultCommitted := false + if reportMatchesProcess && record.ToolCallID != uuid.Nil { + request := agentexecution.ProcessResult{ID: record.ID, Observed: input.Result} + if record.ExecutionSpec.FileTransfer != nil { + _, result, err := processresult.Terminal(record.resultFacts()) + if err != nil { + return DaemonProcessReportApplication{}, err + } + if len(input.Result) > 0 && string(input.Result) != "null" { + result, err = commandTerminalToolResult(record.ID, input.Result) + if err != nil { + return DaemonProcessReportApplication{}, err + } + } + outcome, raw, err := fileTransferToolResultContentParts(ctx, qtx, record, result) + if err != nil { + return DaemonProcessReportApplication{}, err + } + parts, err := parseToolResultContentBlocks(raw) + if err != nil { + return DaemonProcessReportApplication{}, err + } + request.Content, err = executionContent(parts) + if err != nil { + return DaemonProcessReportApplication{}, err + } + request.Outcome = string(outcome) + } + h, err := unit.Handle(input.ProjectID, input.AgentID) + if err != nil { + return DaemonProcessReportApplication{}, err + } + completed, err := h.CompleteProcess(ctx, request) + if err != nil { + return DaemonProcessReportApplication{}, err + } + resultCommitted = completed.Matches + row, err := qtx.GetProcess( + ctx, + dbsqlc.GetProcessParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: record.ID}, + ) + if err != nil { + return DaemonProcessReportApplication{}, err + } + record = processRecordFromSQLC(row) + } else if record.ToolCallID != uuid.Nil { + published, checkErr := publishedToolCallResultExistsTx( + ctx, + qtx, + input.ProjectID, + input.AgentID, + record.ToolCallID, + ) + if checkErr != nil { + return DaemonProcessReportApplication{}, checkErr + } + if !published { + return DaemonProcessReportApplication{}, fmt.Errorf( + "terminal process %s has no durable tool result", + record.ID, + ) + } } - parts, err := parseToolResultContentBlocks(raw) - if err != nil { - return DaemonProcessReportApplication{}, err + if input.StorageExhausted { + if err := completeUnresolvedProcessActionsForClosedProcessTx( + ctx, + unit, + qtx, + record.OrgID, + record.ID, + daemonprotocol.ProcessReasonMachineStorageExhausted, + ); err != nil { + return DaemonProcessReportApplication{}, err + } + } else if processUpdated { + if err := completeQueuedProcessActionsForTerminalProcessTx( + ctx, + unit, + qtx, + record, + ); err != nil { + return DaemonProcessReportApplication{}, err + } } - request.Content, err = executionContent(parts) - if err != nil { + if err := unit.Commit(ctx, "complete daemon process"); err != nil { return DaemonProcessReportApplication{}, err } - request.Outcome = string(outcome) - } - h, err := unit.Handle(input.ProjectID, input.AgentID) - if err != nil { - return DaemonProcessReportApplication{}, err - } - completed, err := h.CompleteProcess(ctx, request) - if err != nil { - return DaemonProcessReportApplication{}, err - } - resultCommitted = completed.Matches - row, err := qtx.GetProcess( - ctx, - dbsqlc.GetProcessParams{ProjectID: input.ProjectID, AgentID: input.AgentID, ID: record.ID}, - ) - if err != nil { - return DaemonProcessReportApplication{}, err - } - record = processRecordFromSQLC(row) - } else if record.ToolCallID != uuid.Nil { - published, checkErr := publishedToolCallResultExistsTx( - ctx, - qtx, - input.ProjectID, - input.AgentID, - record.ToolCallID, - ) - if checkErr != nil { - return DaemonProcessReportApplication{}, checkErr - } - if !published { - return DaemonProcessReportApplication{}, fmt.Errorf( - "terminal process %s has no durable tool result", - record.ID, - ) - } - } - if input.StorageExhausted { - if err := completeUnresolvedProcessActionsForClosedProcessTx( - ctx, - unit, - qtx, - record.OrgID, - record.ID, - daemonprotocol.ProcessReasonMachineStorageExhausted, - ); err != nil { - return DaemonProcessReportApplication{}, err - } - } else if processUpdated { - if err := completeQueuedProcessActionsForTerminalProcessTx( - ctx, - unit, - qtx, - record, - ); err != nil { - return DaemonProcessReportApplication{}, err - } - } - if err := unit.Commit(ctx, "complete daemon process"); err != nil { - return DaemonProcessReportApplication{}, err - } - return DaemonProcessReportApplication{ - Process: record, - ToolResultCommitted: resultCommitted, - }, nil + return DaemonProcessReportApplication{ + Process: record, + ToolResultCommitted: resultCommitted, + }, nil + }) } func daemonProcessForReportTx( diff --git a/internal/storage/executionstore/process_tool_unreachable_store.go b/internal/storage/executionstore/process_tool_unreachable_store.go index dd7d76308..efc6e0b93 100644 --- a/internal/storage/executionstore/process_tool_unreachable_store.go +++ b/internal/storage/executionstore/process_tool_unreachable_store.go @@ -192,68 +192,72 @@ func (s *Store) expireQueuedProcessToolCall( process ProcessRecord, graceSeconds int32, ) (bool, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return false, fmt.Errorf("begin queued process expiry: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - qtx := dbsqlc.New(tx) - unreachable, err := machineStillUnreachableForToolExpiryTx( - ctx, - qtx, - process.OrgID, - process.MachineID, - process.CreatedAt, - graceSeconds, - ) - if err != nil { - return false, err - } - if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: process.ProjectID, - ID: process.AgentID, - }, agentexecution.ExternalAuthority{}); err != nil { - return false, fmt.Errorf("lock agent for queued process expiry: %w", err) - } - row, err := qtx.ExpireQueuedProcessToolCall( - ctx, - dbsqlc.ExpireQueuedProcessToolCallParams{ - ProjectID: process.ProjectID, - AgentID: process.AgentID, - ID: process.ID, - OrgID: process.OrgID, - MachineID: process.MachineID, - QueueTimeoutSeconds: int32(ProcessQueueTimeout / time.Second), - MachineUnreachable: unreachable, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - if err := unit.Commit(ctx, "missed queued process expiry"); err != nil { - return false, err - } - return false, nil - } - if err != nil { - return false, fmt.Errorf("mark queued process failed: %w", err) - } - record := processRecordFromSQLC(row) - h, err := unit.Handle(record.ProjectID, record.AgentID) - if err != nil { - return false, err - } - if record.ToolCallID != uuid.Nil { - if _, - err := h.CompleteProcess(ctx, - agentexecution.ProcessResult{ID: record.ID, - Observed: nil}); err != nil { - return false, err - } - } - if err := unit.Commit(ctx, "queued process expiry"); err != nil { - return false, err - } - return true, nil + return agentexecution.Retry(ctx, + "expireQueuedProcessToolCall", + func(ctx context.Context) (bool, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return false, fmt.Errorf("begin queued process expiry: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) + unreachable, err := machineStillUnreachableForToolExpiryTx( + ctx, + qtx, + process.OrgID, + process.MachineID, + process.CreatedAt, + graceSeconds, + ) + if err != nil { + return false, err + } + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: process.ProjectID, + ID: process.AgentID, + }, nil); err != nil { + return false, fmt.Errorf("lock agent for queued process expiry: %w", err) + } + row, err := qtx.ExpireQueuedProcessToolCall( + ctx, + dbsqlc.ExpireQueuedProcessToolCallParams{ + ProjectID: process.ProjectID, + AgentID: process.AgentID, + ID: process.ID, + OrgID: process.OrgID, + MachineID: process.MachineID, + QueueTimeoutSeconds: int32(ProcessQueueTimeout / time.Second), + MachineUnreachable: unreachable, + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + if err := unit.Commit(ctx, "missed queued process expiry"); err != nil { + return false, err + } + return false, nil + } + if err != nil { + return false, fmt.Errorf("mark queued process failed: %w", err) + } + record := processRecordFromSQLC(row) + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { + return false, err + } + if record.ToolCallID != uuid.Nil { + if _, + err := h.CompleteProcess(ctx, + agentexecution.ProcessResult{ID: record.ID, + Observed: nil}); err != nil { + return false, err + } + } + if err := unit.Commit(ctx, "queued process expiry"); err != nil { + return false, err + } + return true, nil + }) } func (s *Store) FailQueuedProcessAfterWakeFailure( @@ -292,65 +296,69 @@ func (s *Store) failMachineUnreachableQueuedProcessActions( action ProcessActionRecord, graceSeconds int32, ) (int64, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return 0, fmt.Errorf("begin machine-unreachable queued action expiry: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - qtx := dbsqlc.New(tx) - unreachable, err := machineStillUnreachableForToolExpiryTx( - ctx, - qtx, - orgID, - machineID, - action.CreatedAt, - graceSeconds, - ) - if err != nil { - return 0, err - } - if !unreachable { - if err := unit.Commit(ctx, "skipped machine-unreachable queued action expiry"); err != nil { - return 0, err - } - return 0, nil - } - if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ - ProjectID: action.ProjectID, - ID: action.AgentID, - }, agentexecution.ExternalAuthority{}); err != nil { - return 0, fmt.Errorf("lock agent for machine-unreachable queued action expiry: %w", err) - } - rows, err := qtx.MarkQueuedProcessActionsFailedForProcess( - ctx, - dbsqlc.MarkQueuedProcessActionsFailedForProcessParams{ - OrgID: action.OrgID, - ProcessID: action.ProcessID, - ActionID: &action.ID, - StateReasonCode: storeutil.TextFromEmpty(ProcessToolReasonMachineUnreachable), - StateReasonMessage: "", - }, - ) - if err != nil { - return 0, fmt.Errorf("mark machine-unreachable queued actions failed: %w", err) - } - for _, row := range rows { - record := processActionRecordFromSQLC(row) - h, err := unit.Handle(record.ProjectID, record.AgentID) - if err != nil { - return 0, err - } - if record.ToolCallID != uuid.Nil { - if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return agentexecution.Retry(ctx, + "failMachineUnreachableQueuedProcessActions", + func(ctx context.Context) (int64, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return 0, fmt.Errorf("begin machine-unreachable queued action expiry: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) + unreachable, err := machineStillUnreachableForToolExpiryTx( + ctx, + qtx, + orgID, + machineID, + action.CreatedAt, + graceSeconds, + ) + if err != nil { return 0, err } - } - } - if err := unit.Commit(ctx, "machine-unreachable queued action expiry"); err != nil { - return 0, err - } - return int64(len(rows)), nil + if !unreachable { + if err := unit.Commit(ctx, "skipped machine-unreachable queued action expiry"); err != nil { + return 0, err + } + return 0, nil + } + if _, err := unit.LockAgent(ctx, dbsqlc.LockAgentInProjectParams{ + ProjectID: action.ProjectID, + ID: action.AgentID, + }, nil); err != nil { + return 0, fmt.Errorf("lock agent for machine-unreachable queued action expiry: %w", err) + } + rows, err := qtx.MarkQueuedProcessActionsFailedForProcess( + ctx, + dbsqlc.MarkQueuedProcessActionsFailedForProcessParams{ + OrgID: action.OrgID, + ProcessID: action.ProcessID, + ActionID: &action.ID, + StateReasonCode: storeutil.TextFromEmpty(ProcessToolReasonMachineUnreachable), + StateReasonMessage: "", + }, + ) + if err != nil { + return 0, fmt.Errorf("mark machine-unreachable queued actions failed: %w", err) + } + for _, row := range rows { + record := processActionRecordFromSQLC(row) + h, err := unit.Handle(record.ProjectID, record.AgentID) + if err != nil { + return 0, err + } + if record.ToolCallID != uuid.Nil { + if _, err := h.SettleInterruptedAction(ctx, record.ID); err != nil { + return 0, err + } + } + } + if err := unit.Commit(ctx, "machine-unreachable queued action expiry"); err != nil { + return 0, err + } + return int64(len(rows)), nil + }) } func (s *Store) completeMachineUnreachableProcessToolCall( @@ -393,68 +401,72 @@ func (s *Store) completeMachineUnreachableProcessActionToolCall( if err != nil { return false, err } - unit, err := s.cell.Begin(ctx) - if err != nil { - return false, fmt.Errorf( - "begin machine-unreachable process action completion: %w", - err, - ) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + return agentexecution.Retry(ctx, + "completeMachineUnreachableProcessActionToolCall", + func(ctx context.Context) (bool, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return false, fmt.Errorf( + "begin machine-unreachable process action completion: %w", + err, + ) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() - qtx := dbsqlc.New(tx) - unreachable, err := machineStillUnreachableForToolExpiryTx( - ctx, - qtx, - orgID, - machineID, - action.CreatedAt, - graceSeconds, - ) - if err != nil { - return false, err - } - if !unreachable { - if err := unit.Commit(ctx, "skipped machine-unreachable process action completion"); err != nil { - return false, err - } - return false, nil - } - state := ProcessActionStateUnknown - if action.ActionKind == ProcessActionKindRead { - state = ProcessActionStateFailed - } - _, err = completeDaemonProcessActionTx( - ctx, - unit, - qtx, - processActionCompletionInput{ - ProjectID: action.ProjectID, - AgentID: action.AgentID, - ProcessID: action.ProcessID, - ID: action.ID, - StateReasonCode: ProcessToolReasonMachineUnreachable, - Result: result, - }, - state, - ) - if errors.Is(err, storeerr.ErrDaemonRuntimeUnregistered) { - if err := unit.Commit(ctx, "missed machine-unreachable process action completion"); err != nil { - return false, err - } - return false, nil - } - if err != nil { - return false, fmt.Errorf( - "resolve machine-unreachable process action: %w", - err, - ) - } - if err := unit.Commit(ctx, "machine-unreachable process action completion"); err != nil { - return false, err - } - return true, nil + qtx := dbsqlc.New(tx) + unreachable, err := machineStillUnreachableForToolExpiryTx( + ctx, + qtx, + orgID, + machineID, + action.CreatedAt, + graceSeconds, + ) + if err != nil { + return false, err + } + if !unreachable { + if err := unit.Commit(ctx, "skipped machine-unreachable process action completion"); err != nil { + return false, err + } + return false, nil + } + state := ProcessActionStateUnknown + if action.ActionKind == ProcessActionKindRead { + state = ProcessActionStateFailed + } + _, err = completeDaemonProcessActionTx( + ctx, + unit, + qtx, + processActionCompletionInput{ + ProjectID: action.ProjectID, + AgentID: action.AgentID, + ProcessID: action.ProcessID, + ID: action.ID, + StateReasonCode: ProcessToolReasonMachineUnreachable, + Result: result, + }, + state, + ) + if errors.Is(err, storeerr.ErrDaemonRuntimeUnregistered) { + if err := unit.Commit(ctx, "missed machine-unreachable process action completion"); err != nil { + return false, err + } + return false, nil + } + if err != nil { + return false, fmt.Errorf( + "resolve machine-unreachable process action: %w", + err, + ) + } + if err := unit.Commit(ctx, "machine-unreachable process action completion"); err != nil { + return false, err + } + return true, nil + }) } func (s *Store) completeMachineUnreachableToolCall( @@ -469,64 +481,68 @@ func (s *Store) completeMachineUnreachableToolCall( if err != nil { return false, err } - unit, err := s.cell.Begin(ctx) - if err != nil { - return false, err - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - qtx := dbsqlc.New(tx) - unreachable, err := machineStillUnreachableForToolExpiryTx( - ctx, - qtx, - orgID, - machineID, - fallbackAt, - graceSeconds, - ) - if err != nil { - return false, err - } - if !unreachable { - if err := unit.Commit(ctx, "skipped machine-unreachable tool call completion"); err != nil { - return false, err - } - return false, nil - } - if _, - err := unit.LockAgent(ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: projectID, - ID: agentID}, - agentexecution.ExternalAuthority{}); err != nil { - return false, err - } - process, found, err := getProcessByToolCallTx(ctx, tx, projectID, agentID, toolCallID) - if err != nil { - return false, err - } - if !found { - return false, storeerr.ErrNotFound - } - blocks, err := parseToolResultContentBlocks(parts) - if err != nil { - return false, err - } - content, err := executionContent(blocks) - if err != nil { - return false, err - } - h, err := unit.Handle(projectID, agentID) - if err != nil { - return false, err - } - completed, err := h.FailProcessTool(ctx, process.ID, content) - if err != nil { - return false, err - } - if err := unit.Commit(ctx, "machine-unreachable tool call completion"); err != nil { - return false, err - } - return completed.Changed, nil + return agentexecution.Retry(ctx, + "completeMachineUnreachableToolCall", + func(ctx context.Context) (bool, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return false, err + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) + unreachable, err := machineStillUnreachableForToolExpiryTx( + ctx, + qtx, + orgID, + machineID, + fallbackAt, + graceSeconds, + ) + if err != nil { + return false, err + } + if !unreachable { + if err := unit.Commit(ctx, "skipped machine-unreachable tool call completion"); err != nil { + return false, err + } + return false, nil + } + if _, + err := unit.LockAgent(ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: projectID, + ID: agentID}, + nil); err != nil { + return false, err + } + process, found, err := getProcessByToolCallTx(ctx, tx, projectID, agentID, toolCallID) + if err != nil { + return false, err + } + if !found { + return false, storeerr.ErrNotFound + } + blocks, err := parseToolResultContentBlocks(parts) + if err != nil { + return false, err + } + content, err := executionContent(blocks) + if err != nil { + return false, err + } + h, err := unit.Handle(projectID, agentID) + if err != nil { + return false, err + } + completed, err := h.FailProcessTool(ctx, process.ID, content) + if err != nil { + return false, err + } + if err := unit.Commit(ctx, "machine-unreachable tool call completion"); err != nil { + return false, err + } + return completed.Changed, nil + }) } func machineStillUnreachableForToolExpiryTx( diff --git a/internal/storage/executionstore/processes_store.go b/internal/storage/executionstore/processes_store.go index 679773045..3256c535e 100644 --- a/internal/storage/executionstore/processes_store.go +++ b/internal/storage/executionstore/processes_store.go @@ -329,50 +329,54 @@ func (s *Store) CompleteProcess(ctx context.Context, input CompleteProcessInput) return ProcessRecord{}, errors.New("source ended at is required") } input.SourceEndedAt = canonicalSourceTime(input.SourceEndedAt) - unit, err := s.cell.Begin(ctx) - if err != nil { - return ProcessRecord{}, fmt.Errorf("begin complete process: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() + return agentexecution.Retry(ctx, + "CompleteProcess", + func(ctx context.Context) (ProcessRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return ProcessRecord{}, fmt.Errorf("begin complete process: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() - qtx := dbsqlc.New(tx) - if _, err := unit.LockAgent( - ctx, - dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, - agentexecution.ExternalAuthority{}, - ); err != nil { - return ProcessRecord{}, fmt.Errorf("lock agent for process completion: %w", err) - } - row, err := qtx.CompleteProcess( - ctx, - dbsqlc.CompleteProcessParams{ - ProjectID: input.ProjectID, - AgentID: input.AgentID, - ID: input.ID, - RuntimeLockID: input.RuntimeLockID, - State: string(input.State), - SourceEndedAt: input.SourceEndedAt, - ExitCode: storeutil.Int32Ptr(input.ExitCode), - ExitSignal: input.ExitSignal, - StateReasonCode: storeutil.TextFromEmpty(input.StateReasonCode), - StateReasonMessage: input.StateReasonMessage, - }, - ) - if err != nil { - return ProcessRecord{}, fmt.Errorf("complete process: %w", err) - } - process := processRecordFromCompleteSQLC(row) - if err := completeQueuedProcessActionsForTerminalProcessTx( - ctx, - unit, - qtx, - process, - ); err != nil { - return ProcessRecord{}, err - } - if err := unit.Commit(ctx, "complete process"); err != nil { - return ProcessRecord{}, err - } - return process, nil + qtx := dbsqlc.New(tx) + if _, err := unit.LockAgent( + ctx, + dbsqlc.LockAgentInProjectParams{ProjectID: input.ProjectID, ID: input.AgentID}, + nil, + ); err != nil { + return ProcessRecord{}, fmt.Errorf("lock agent for process completion: %w", err) + } + row, err := qtx.CompleteProcess( + ctx, + dbsqlc.CompleteProcessParams{ + ProjectID: input.ProjectID, + AgentID: input.AgentID, + ID: input.ID, + RuntimeLockID: input.RuntimeLockID, + State: string(input.State), + SourceEndedAt: input.SourceEndedAt, + ExitCode: storeutil.Int32Ptr(input.ExitCode), + ExitSignal: input.ExitSignal, + StateReasonCode: storeutil.TextFromEmpty(input.StateReasonCode), + StateReasonMessage: input.StateReasonMessage, + }, + ) + if err != nil { + return ProcessRecord{}, fmt.Errorf("complete process: %w", err) + } + process := processRecordFromCompleteSQLC(row) + if err := completeQueuedProcessActionsForTerminalProcessTx( + ctx, + unit, + qtx, + process, + ); err != nil { + return ProcessRecord{}, err + } + if err := unit.Commit(ctx, "complete process"); err != nil { + return ProcessRecord{}, err + } + return process, nil + }) } diff --git a/internal/storage/executionstore/project_machine_pool_grants_store.go b/internal/storage/executionstore/project_machine_pool_grants_store.go index b44aa50d6..ec939dd1c 100644 --- a/internal/storage/executionstore/project_machine_pool_grants_store.go +++ b/internal/storage/executionstore/project_machine_pool_grants_store.go @@ -762,8 +762,8 @@ func (s *Store) DeleteProjectMachinePoolGrant( if orgID == uuid.Nil || projectID == uuid.Nil || id == uuid.Nil { return DeleteProjectMachinePoolGrantResult{}, errors.New("pool grant org, project, and id are required") } - return storeutil.RetryTransaction(ctx, "delete_project_machine_pool_grant", - func() (DeleteProjectMachinePoolGrantResult, error) { + return agentexecution.Retry(ctx, "delete_project_machine_pool_grant", + func(ctx context.Context) (DeleteProjectMachinePoolGrantResult, error) { return s.deleteProjectMachinePoolGrantOnce(ctx, orgID, projectID, id) }) } @@ -857,7 +857,7 @@ func (s *Store) deleteProjectMachinePoolGrantOnce( AgentID: agentRow.AgentID, }) } - if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, nil); err != nil { return DeleteProjectMachinePoolGrantResult{}, err } machineRows, err := qtx.MarkPoolGrantMachinesDeleting( diff --git a/internal/storage/executionstore/provider_runtime_reconciliation_store.go b/internal/storage/executionstore/provider_runtime_reconciliation_store.go index 5f5af1b76..6c4a8498d 100644 --- a/internal/storage/executionstore/provider_runtime_reconciliation_store.go +++ b/internal/storage/executionstore/provider_runtime_reconciliation_store.go @@ -9,6 +9,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" @@ -348,68 +349,77 @@ func (s *Store) ClaimProviderRuntimeMismatchDeletion( if !providerRuntimeWakeAllowsDeletion(candidate) { return PoolMachineDeletionClaim{}, false, nil } - unit, err := s.cell.Begin(ctx) - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf( - "begin provider runtime mismatch deletion claim: %w", - err, - ) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - txNotifications := unit.Notifications() - qtx := dbsqlc.New(tx) - locked, err := lockProviderRuntimeDeletionCandidate(ctx, tx, qtx, candidate) - if err != nil { - return PoolMachineDeletionClaim{}, false, err - } - if !locked { - return PoolMachineDeletionClaim{}, false, nil - } - row, err := qtx.ClaimProviderRuntimeMismatchDeletion( - ctx, - dbsqlc.ClaimProviderRuntimeMismatchDeletionParams{ - ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), - OrgID: candidate.OrgID, - MachineID: candidate.MachineID, - MachinePoolID: candidate.MachinePoolID, - LifecycleVersion: candidate.LifecycleVersion, - Provider: candidate.Provider, - ProviderResourceID: storeutil.TextFromEmpty(candidate.ProviderResourceID), - MismatchSince: *candidate.ProviderRuntimeMismatchSince, - ConfirmationGraceMilliseconds: input.ConfirmationGrace.Milliseconds(), - WakeAttemptExpiresAt: candidate.WakeAttemptExpiresAt, - DaemonRuntimeID: candidate.CurrentDaemonRuntimeID, - InactiveSince: candidate.InactiveSince, - InactivityGraceMilliseconds: input.InactivityGrace.Milliseconds(), - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return PoolMachineDeletionClaim{}, false, nil - } - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf( - "claim provider runtime mismatch deletion: %w", - err, - ) - } - claim := providerRuntimeMismatchDeletionClaimFromSQLC(row) - claim, err = s.finalizePoolMachineDeletionClaimTx( - ctx, - unit, - qtx, - txNotifications, - claim, + type transactionResult struct { + claim PoolMachineDeletionClaim + claimed bool + } + retried, retryErr := agentexecution.Retry(ctx, + "ClaimProviderRuntimeMismatchDeletion", + func(ctx context.Context) (transactionResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return transactionResult{}, fmt.Errorf( + "begin provider runtime mismatch deletion claim: %w", + err, + ) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + txNotifications := unit.Notifications() + qtx := dbsqlc.New(tx) + locked, err := lockProviderRuntimeDeletionCandidate(ctx, tx, qtx, candidate) + if err != nil { + return transactionResult{}, err + } + if !locked { + return transactionResult{}, nil + } + row, err := qtx.ClaimProviderRuntimeMismatchDeletion( + ctx, + dbsqlc.ClaimProviderRuntimeMismatchDeletionParams{ + ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), + OrgID: candidate.OrgID, + MachineID: candidate.MachineID, + MachinePoolID: candidate.MachinePoolID, + LifecycleVersion: candidate.LifecycleVersion, + Provider: candidate.Provider, + ProviderResourceID: storeutil.TextFromEmpty(candidate.ProviderResourceID), + MismatchSince: *candidate.ProviderRuntimeMismatchSince, + ConfirmationGraceMilliseconds: input.ConfirmationGrace.Milliseconds(), + WakeAttemptExpiresAt: candidate.WakeAttemptExpiresAt, + DaemonRuntimeID: candidate.CurrentDaemonRuntimeID, + InactiveSince: candidate.InactiveSince, + InactivityGraceMilliseconds: input.InactivityGrace.Milliseconds(), + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return transactionResult{}, nil + } + if err != nil { + return transactionResult{}, fmt.Errorf( + "claim provider runtime mismatch deletion: %w", + err, + ) + } + claim := providerRuntimeMismatchDeletionClaimFromSQLC(row) + claim, err = s.finalizePoolMachineDeletionClaimTx( + ctx, + unit, + qtx, + txNotifications, + claim, - ProcessToolReasonMachineUnreachable, - ) - if err != nil { - return PoolMachineDeletionClaim{}, false, err - } - if err := unit.Commit(ctx, "claim provider runtime mismatch deletion"); err != nil { - return PoolMachineDeletionClaim{}, false, err - } - return claim, true, nil + ProcessToolReasonMachineUnreachable, + ) + if err != nil { + return transactionResult{}, err + } + if err := unit.Commit(ctx, "claim provider runtime mismatch deletion"); err != nil { + return transactionResult{}, err + } + return transactionResult{claim, true}, nil + }) + return retried.claim, retried.claimed, retryErr } func (s *Store) ClaimProviderRuntimeTerminatedDeletion( @@ -424,64 +434,73 @@ func (s *Store) ClaimProviderRuntimeTerminatedDeletion( if !providerRuntimeWakeAllowsDeletion(candidate) { return PoolMachineDeletionClaim{}, false, nil } - unit, err := s.cell.Begin(ctx) - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf( - "begin provider runtime terminated deletion claim: %w", - err, - ) - } - defer func() { _ = unit.Rollback(ctx) }() - tx := unit.DB() - qtx := dbsqlc.New(tx) - locked, err := lockProviderRuntimeDeletionCandidate(ctx, tx, qtx, candidate) - if err != nil { - return PoolMachineDeletionClaim{}, false, err - } - if !locked { - return PoolMachineDeletionClaim{}, false, nil - } - row, err := qtx.ClaimProviderRuntimeTerminatedDeletion( - ctx, - dbsqlc.ClaimProviderRuntimeTerminatedDeletionParams{ - ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), - OrgID: candidate.OrgID, - MachineID: candidate.MachineID, - MachinePoolID: candidate.MachinePoolID, - LifecycleVersion: candidate.LifecycleVersion, - Provider: candidate.Provider, - ProviderResourceID: storeutil.TextFromEmpty(candidate.ProviderResourceID), - MismatchSince: candidate.ProviderRuntimeMismatchSince, - WakeAttemptExpiresAt: candidate.WakeAttemptExpiresAt, - DaemonRuntimeID: candidate.CurrentDaemonRuntimeID, - InactiveSince: candidate.InactiveSince, - }, - ) - if errors.Is(err, pgx.ErrNoRows) { - return PoolMachineDeletionClaim{}, false, nil - } - if err != nil { - return PoolMachineDeletionClaim{}, false, fmt.Errorf( - "claim provider runtime terminated deletion: %w", - err, - ) - } - claim := providerRuntimeTerminatedDeletionClaimFromSQLC(row) - claim, err = s.finalizePoolMachineDeletionClaimTx( - ctx, - unit, - qtx, - unit.Notifications(), - claim, - ProcessToolReasonMachineUnreachable, - ) - if err != nil { - return PoolMachineDeletionClaim{}, false, err - } - if err := unit.Commit(ctx, "claim provider runtime terminated deletion"); err != nil { - return PoolMachineDeletionClaim{}, false, err - } - return claim, true, nil + type transactionResult struct { + claim PoolMachineDeletionClaim + claimed bool + } + retried, retryErr := agentexecution.Retry(ctx, + "ClaimProviderRuntimeTerminatedDeletion", + func(ctx context.Context) (transactionResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return transactionResult{}, fmt.Errorf( + "begin provider runtime terminated deletion claim: %w", + err, + ) + } + defer func() { _ = unit.Rollback(ctx) }() + tx := unit.DB() + qtx := dbsqlc.New(tx) + locked, err := lockProviderRuntimeDeletionCandidate(ctx, tx, qtx, candidate) + if err != nil { + return transactionResult{}, err + } + if !locked { + return transactionResult{}, nil + } + row, err := qtx.ClaimProviderRuntimeTerminatedDeletion( + ctx, + dbsqlc.ClaimProviderRuntimeTerminatedDeletionParams{ + ClaimTimeoutSeconds: int64(poolMachineDeletionLeaseDuration / time.Second), + OrgID: candidate.OrgID, + MachineID: candidate.MachineID, + MachinePoolID: candidate.MachinePoolID, + LifecycleVersion: candidate.LifecycleVersion, + Provider: candidate.Provider, + ProviderResourceID: storeutil.TextFromEmpty(candidate.ProviderResourceID), + MismatchSince: candidate.ProviderRuntimeMismatchSince, + WakeAttemptExpiresAt: candidate.WakeAttemptExpiresAt, + DaemonRuntimeID: candidate.CurrentDaemonRuntimeID, + InactiveSince: candidate.InactiveSince, + }, + ) + if errors.Is(err, pgx.ErrNoRows) { + return transactionResult{}, nil + } + if err != nil { + return transactionResult{}, fmt.Errorf( + "claim provider runtime terminated deletion: %w", + err, + ) + } + claim := providerRuntimeTerminatedDeletionClaimFromSQLC(row) + claim, err = s.finalizePoolMachineDeletionClaimTx( + ctx, + unit, + qtx, + unit.Notifications(), + claim, + ProcessToolReasonMachineUnreachable, + ) + if err != nil { + return transactionResult{}, err + } + if err := unit.Commit(ctx, "claim provider runtime terminated deletion"); err != nil { + return transactionResult{}, err + } + return transactionResult{claim, true}, nil + }) + return retried.claim, retried.claimed, retryErr } func providerRuntimeDeletionCandidateValid(candidate ProviderRuntimeCandidate) bool { diff --git a/internal/storage/executionstore/store.go b/internal/storage/executionstore/store.go index 11e8d370d..c3d03d28a 100644 --- a/internal/storage/executionstore/store.go +++ b/internal/storage/executionstore/store.go @@ -39,7 +39,6 @@ func New(pool *pgxpool.Pool, config Config) *Store { db := storeutil.WrapPool(pool) return &Store{ cell: agentexecution.NewCell( - "primary", pool, config.PostCommitPublisher, config.ModelCallRetryBackoff, diff --git a/internal/storage/executionstore/subagents_store.go b/internal/storage/executionstore/subagents_store.go index f79267f5c..802b31469 100644 --- a/internal/storage/executionstore/subagents_store.go +++ b/internal/storage/executionstore/subagents_store.go @@ -160,7 +160,7 @@ func admitSubagentLaunchTx( if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ ProjectID: projectID, AgentID: launch.ParentAgentID, - }}, agentexecution.LifecycleAuthority{}); err != nil { + }}, nil); err != nil { if errors.Is(err, storeerr.ErrNotFound) { return fmt.Errorf("parent agent: %w", storeerr.ErrNotFound) } @@ -441,7 +441,7 @@ func (t *toolCallTransaction) sendSubagentMessage(ctx context.Context, input Sen {ProjectID: t.input.ProjectID, AgentID: t.input.AgentID}, {ProjectID: t.input.ProjectID, AgentID: input.TargetAgentID}, }, - agentexecution.RuntimeAuthority{AgentID: t.input.AgentID, + &agentexecution.RuntimeAuthority{AgentID: t.input.AgentID, RuntimeLockID: t.input.RuntimeLockID}); err != nil { return err } @@ -581,7 +581,7 @@ func (t *toolCallTransaction) cancelSubagent(ctx context.Context, child AgentRec {ProjectID: child.ProjectID, AgentID: child.ID}, {ProjectID: t.input.ProjectID, AgentID: t.input.AgentID}, }, - agentexecution.RuntimeAuthority{AgentID: t.input.AgentID, + &agentexecution.RuntimeAuthority{AgentID: t.input.AgentID, RuntimeLockID: t.input.RuntimeLockID}); err != nil { return err } @@ -646,10 +646,10 @@ func (s *Store) archiveIdleCandidates( var machines []MachineRecord archived := 0 for _, candidate := range candidates { - released, err := storeutil.RetryTransaction( + released, err := agentexecution.Retry( ctx, "archive_idle_agent", - func() ([]MachineRecord, error) { + func(ctx context.Context) ([]MachineRecord, error) { return s.archiveIdleCandidateOnce(ctx, candidate, asOf) }, ) diff --git a/internal/storage/executionstore/tool_call_execution.go b/internal/storage/executionstore/tool_call_execution.go index 3827802f0..e481ea00e 100644 --- a/internal/storage/executionstore/tool_call_execution.go +++ b/internal/storage/executionstore/tool_call_execution.go @@ -11,7 +11,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -77,9 +76,11 @@ func (s *Store) ExecuteToolCall( if plan == nil { return ExecuteToolCallResult{}, errors.New("tool call plan is required") } - return storeutil.RetryTransaction(ctx, "execute_tool_call", func() (ExecuteToolCallResult, error) { - return s.executeToolCallOnce(ctx, input, plan) - }) + return agentexecution.Retry(ctx, + "execute_tool_call", + func(ctx context.Context) (ExecuteToolCallResult, error) { + return s.executeToolCallOnce(ctx, input, plan) + }) } func (s *Store) executeToolCallOnce( @@ -162,7 +163,7 @@ func (t *toolCallTransaction) lockToolCall(ctx context.Context, acceptExisting b } err := t.unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: t.input.ProjectID, AgentID: t.input.AgentID}}, - agentexecution.IngressAuthority{}, + nil, ) if errors.Is(err, storeerr.ErrNotFound) { return storeerr.ErrRuntimeLockInactive diff --git a/internal/storage/executionstore/tool_calls_store.go b/internal/storage/executionstore/tool_calls_store.go index 81d2d07c5..1a70e9af3 100644 --- a/internal/storage/executionstore/tool_calls_store.go +++ b/internal/storage/executionstore/tool_calls_store.go @@ -146,7 +146,7 @@ func (s *Store) CompleteToolCall( return ToolCallRecord{}, err } input.ResultContentParts = parts - record, err := s.completeToolCallOnce(ctx, input) + record, err := s.completePreparedToolCall(ctx, input) if err == nil { return record, nil } @@ -155,64 +155,77 @@ func (s *Store) CompleteToolCall( return ToolCallRecord{}, err } input.ResultContentParts = parts - return s.completeToolCallOnce(ctx, input) + return s.completePreparedToolCall(ctx, input) } -func (s *Store) completeToolCallOnce( +func (s *Store) completePreparedToolCall( ctx context.Context, input CompleteToolCallInput, ) (ToolCallRecord, error) { - unit, err := s.cell.Begin(ctx) - if err != nil { - return ToolCallRecord{}, fmt.Errorf("begin complete tool call: %w", err) - } - defer func() { _ = unit.Rollback(ctx) }() - - record, err := completeExecutionTool(ctx, unit, input) - if err != nil { - return ToolCallRecord{}, err - } - if err := unit.Commit(ctx, "complete tool call"); err != nil { - return ToolCallRecord{}, err - } - return record, nil + return agentexecution.Retry(ctx, + "completePreparedToolCall", + func(ctx context.Context) (ToolCallRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return ToolCallRecord{}, fmt.Errorf("begin complete tool call: %w", err) + } + defer func() { _ = unit.Rollback(ctx) }() + + record, err := completeExecutionTool(ctx, unit, input) + if err != nil { + return ToolCallRecord{}, err + } + if err := unit.Commit(ctx, "complete tool call"); err != nil { + return ToolCallRecord{}, err + } + return record, nil + }) } func (s *Store) MarkToolCallReady(ctx context.Context, input MarkToolCallReadyInput) (ToolCallRecord, error) { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return ToolCallRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - if _, - err = h.AuthorizeTool(ctx, - agentexecution.ToolRef{ID: input.ID, - RuntimeLockID: input.RuntimeLockID}); err != nil { - return ToolCallRecord{}, err - } - record, err := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, input.ID) - if err != nil { - return ToolCallRecord{}, err - } - if err = unit.Commit(ctx, "authorize tool"); err != nil { - return ToolCallRecord{}, err - } - return record, nil + return agentexecution.Retry(ctx, + "MarkToolCallReady", + func(ctx context.Context) (ToolCallRecord, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return ToolCallRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if _, + err = h.AuthorizeTool(ctx, + agentexecution.ToolRef{ID: input.ID, + RuntimeLockID: input.RuntimeLockID}); err != nil { + return ToolCallRecord{}, err + } + record, err := getToolCallTx(ctx, unit.DB(), input.ProjectID, input.AgentID, input.ID) + if err != nil { + return ToolCallRecord{}, err + } + if err = unit.Commit(ctx, "authorize tool"); err != nil { + return ToolCallRecord{}, err + } + return record, nil + }) } func (s *Store) RequeueRuntimeToolCall(ctx context.Context, input RequeueRuntimeToolCallInput) error { - unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) - if err != nil { - return err - } - defer func() { _ = unit.Rollback(ctx) }() - if _, - err = h.RequeueTool(ctx, - agentexecution.ToolRef{ID: input.ToolCallID, - RuntimeLockID: input.RuntimeLockID}); err != nil { - return err - } - return unit.Commit(ctx, "requeue tool") + _, err := agentexecution.Retry(ctx, + "RequeueRuntimeToolCall", + func(ctx context.Context) (struct{}, error) { + unit, h, err := beginExecution(ctx, s, input.ProjectID, input.AgentID) + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if _, + err = h.RequeueTool(ctx, + agentexecution.ToolRef{ID: input.ToolCallID, + RuntimeLockID: input.RuntimeLockID}); err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "requeue tool") + }) + return err } func completeExecutionTool( @@ -223,7 +236,7 @@ func completeExecutionTool( err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, AgentID: input.AgentID}}, - agentexecution.IngressAuthority{}, + nil, ) if errors.Is(err, storeerr.ErrNotFound) { return ToolCallRecord{}, storeerr.ErrRuntimeLockInactive @@ -261,7 +274,7 @@ func (s *Store) CompleteRuntimeToolCall( return ToolCallRecord{}, err } input.ResultContentParts = parts - record, err := s.completeRuntimeToolCallOnce(ctx, input) + record, err := s.completePreparedRuntimeToolCall(ctx, input) if err == nil { return record, nil } @@ -270,37 +283,41 @@ func (s *Store) CompleteRuntimeToolCall( return ToolCallRecord{}, err } input.ResultContentParts = parts - return s.completeRuntimeToolCallOnce(ctx, input) + return s.completePreparedRuntimeToolCall(ctx, input) } -func (s *Store) completeRuntimeToolCallOnce( +func (s *Store) completePreparedRuntimeToolCall( ctx context.Context, input CompleteRuntimeToolCallInput, ) (ToolCallRecord, error) { if !input.Outcome.IsTerminal() || input.Outcome == ToolResultOutcomeDenied { return ToolCallRecord{}, errors.New("invalid runtime tool outcome") } - unit, err := s.cell.Begin(ctx) - if err != nil { - return ToolCallRecord{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - record, err := completeExecutionTool( - ctx, - unit, - CompleteToolCallInput{ProjectID: input.ProjectID, AgentID: input.AgentID, - ID: input.ID, - RuntimeLockID: input.RuntimeLockID, - Outcome: input.Outcome, - ResultContentParts: input.ResultContentParts}, - ) - if err != nil { - return ToolCallRecord{}, err - } - if err = unit.Commit(ctx, "complete runtime tool"); err != nil { - return ToolCallRecord{}, err - } - return record, nil + return agentexecution.Retry(ctx, + "completePreparedRuntimeToolCall", + func(ctx context.Context) (ToolCallRecord, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return ToolCallRecord{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + record, err := completeExecutionTool( + ctx, + unit, + CompleteToolCallInput{ProjectID: input.ProjectID, AgentID: input.AgentID, + ID: input.ID, + RuntimeLockID: input.RuntimeLockID, + Outcome: input.Outcome, + ResultContentParts: input.ResultContentParts}, + ) + if err != nil { + return ToolCallRecord{}, err + } + if err = unit.Commit(ctx, "complete runtime tool"); err != nil { + return ToolCallRecord{}, err + } + return record, nil + }) } func (s *Store) CompleteCustomToolCall( @@ -312,20 +329,24 @@ func (s *Store) CompleteCustomToolCall( return CompleteCustomToolCallResult{}, err } input.ContentBlocks = parts - unit, err := s.cell.Begin(ctx) - if err != nil { - return CompleteCustomToolCallResult{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - - result, err := completeCustomToolCallTx(ctx, unit, input) - if err != nil { - return CompleteCustomToolCallResult{}, err - } - if err := unit.Commit(ctx, "complete custom tool call"); err != nil { - return CompleteCustomToolCallResult{}, err - } - return result, nil + return agentexecution.Retry(ctx, + "CompleteCustomToolCall", + func(ctx context.Context) (CompleteCustomToolCallResult, error) { + unit, err := s.cell.Begin(ctx) + if err != nil { + return CompleteCustomToolCallResult{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + + result, err := completeCustomToolCallTx(ctx, unit, input) + if err != nil { + return CompleteCustomToolCallResult{}, err + } + if err := unit.Commit(ctx, "complete custom tool call"); err != nil { + return CompleteCustomToolCallResult{}, err + } + return result, nil + }) } func completeCustomToolCallTx( @@ -336,7 +357,7 @@ func completeCustomToolCallTx( if err := unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: input.ProjectID, AgentID: input.AgentID}}, - agentexecution.ExternalAuthority{}); err != nil { + nil); err != nil { return CompleteCustomToolCallResult{}, err } blocks, err := parseToolResultContentBlocks(input.ContentBlocks) diff --git a/internal/storage/integrationstore/export_integration.go b/internal/storage/integrationstore/export_integration.go index 4aaef4b8a..299c88a3c 100644 --- a/internal/storage/integrationstore/export_integration.go +++ b/internal/storage/integrationstore/export_integration.go @@ -6,6 +6,7 @@ import ( "context" "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) func (s *Store) DeleteIntegrationOnceForIntegration( @@ -16,5 +17,8 @@ func (s *Store) DeleteIntegrationOnceForIntegration( if err != nil { return err } - return s.deleteIntegrationOnce(ctx, integration.OrgID, projectID, integrationID) + _, err = agentexecution.Retry(ctx, "test integration deletion", func(ctx context.Context) (struct{}, error) { + return struct{}{}, s.deleteIntegrationOnce(ctx, integration.OrgID, projectID, integrationID) + }) + return err } diff --git a/internal/storage/integrationstore/integration_lifecycle.go b/internal/storage/integrationstore/integration_lifecycle.go index 4631a2b44..c0fb1ce0f 100644 --- a/internal/storage/integrationstore/integration_lifecycle.go +++ b/internal/storage/integrationstore/integration_lifecycle.go @@ -10,7 +10,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -67,7 +66,7 @@ func (s *Store) DeleteIntegration(ctx context.Context, orgID, projectID, id uuid if orgID == uuid.Nil || projectID == uuid.Nil || id == uuid.Nil { return storeerr.InvalidRequest(errors.New("organization, project and integration are required")) } - _, err := storeutil.RetryTransaction(ctx, "delete_integration", func() (struct{}, error) { + _, err := agentexecution.Retry(ctx, "delete_integration", func(ctx context.Context) (struct{}, error) { return struct{}{}, s.deleteIntegrationOnce(ctx, orgID, projectID, id) }) return err @@ -104,7 +103,7 @@ func (s *Store) deleteIntegrationOnce(ctx context.Context, orgID, projectID, id for _, agentID := range ids { refs = append(refs, lifecyclelock.AgentRef{ProjectID: projectID, AgentID: agentID}) } - if err := unit.LockAgentRefs(ctx, refs, agentexecution.LifecycleAuthority{}); err != nil { + if err := unit.LockAgentRefs(ctx, refs, nil); err != nil { return err } if err := q.DeleteIntegrationSubscriptions(ctx, dbsqlc.DeleteIntegrationSubscriptionsParams{ diff --git a/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go b/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go index 9f1f2ee4f..b9f83c29f 100644 --- a/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go +++ b/internal/storage/integrationstore/integration_target_lifecycle_integration_test.go @@ -52,7 +52,7 @@ FROM identities identity CROSS JOIN agents agent WHERE agent.id=$1 RETURNING id) INSERT INTO agent_execution_state(agent_id,turn_continuable,incomplete_tools) SELECT id,false,false FROM inserted`, current.ID, ) - cell := agentexecution.NewCell("test", f.pool, nil, nil) + cell := agentexecution.NewCell(f.pool, nil, nil) require.NoError(t, cell.Transact(f.ctx, func(u *agentexecution.Unit) error { return u.ClearIntegrationTargets(f.ctx, uuid.New(), f.integrationID) })) diff --git a/internal/storage/integrationstore/store.go b/internal/storage/integrationstore/store.go index 76031dd73..eaf728504 100644 --- a/internal/storage/integrationstore/store.go +++ b/internal/storage/integrationstore/store.go @@ -32,7 +32,7 @@ type Store struct { func New(pool *pgxpool.Pool, config Config) *Store { db := storeutil.WrapPool(pool) return &Store{ - cell: agentexecution.NewCell("primary", pool, config.PostCommitPublisher, nil), + cell: agentexecution.NewCell(pool, config.PostCommitPublisher, nil), pool: db, q: dbsqlc.New(db), access: config.Access, diff --git a/internal/storage/internal/agentexecution/cell.go b/internal/storage/internal/agentexecution/cell.go index c33dcbe27..f4bd3d309 100644 --- a/internal/storage/internal/agentexecution/cell.go +++ b/internal/storage/internal/agentexecution/cell.go @@ -14,14 +14,12 @@ import ( ) type AgentCell struct { - id CellID pool *storeutil.Pool publisher notifications.PostCommitPublisher retryBackoff func(int, string) time.Duration } func NewCell( - id CellID, pool *pgxpool.Pool, publisher notifications.PostCommitPublisher, retryBackoff func(int, string) time.Duration, @@ -30,34 +28,37 @@ func NewCell( retryBackoff = modelprotocol.RetryBackoff } return &AgentCell{ - id: id, pool: storeutil.WrapPool(pool), publisher: publisher, retryBackoff: retryBackoff, } } -func (c *AgentCell) ID() CellID { return c.id } - func (c *AgentCell) Begin(ctx context.Context) (*Unit, error) { - if c.id == "" { - return nil, errors.New("agent cell identity is required") + plan, ok := ctx.Value(lockPlanKey{}).(retryPlan) + if !ok { + return nil, errors.New("agent execution unit requires a retry boundary") } tx, err := c.pool.BeginTx(ctx, pgx.TxOptions{}) if err != nil { return nil, err } - return newUnit(c, tx), nil + u := newUnit(c, tx) + u.expanded = plan + return u, nil } func (c *AgentCell) Transact(ctx context.Context, run func(*Unit) error) error { - u, err := c.Begin(ctx) - if err != nil { - return err - } - defer func() { _ = u.Rollback(context.WithoutCancel(ctx)) }() - if err := run(u); err != nil { - return err - } - return u.Commit(ctx, "agent execution") + _, err := Retry(ctx, "agent execution", func(ctx context.Context) (struct{}, error) { + u, err := c.Begin(ctx) + if err != nil { + return struct{}{}, err + } + defer func() { _ = u.Rollback(context.WithoutCancel(ctx)) }() + if err = run(u); err != nil { + return struct{}{}, err + } + return struct{}{}, u.Commit(ctx, "agent execution") + }) + return err } diff --git a/internal/storage/internal/agentexecution/commands_agents.go b/internal/storage/internal/agentexecution/commands_agents.go index 2b004da22..f91a2b859 100644 --- a/internal/storage/internal/agentexecution/commands_agents.go +++ b/internal/storage/internal/agentexecution/commands_agents.go @@ -122,7 +122,8 @@ func (h *Handle) parentEffect(ctx context.Context, kind, text, key string, inter } parent := h.unit.handles[valueOrZero(child.ParentAgentID)] if parent == nil { - return ErrLockPlan + return h.unit.expandPlan(lifecyclelock.AgentRef{ProjectID: h.route.ProjectID, + AgentID: valueOrZero(child.ParentAgentID)}) } snapshot, err := parent.LoadExecution(ctx) if err != nil { @@ -239,7 +240,7 @@ func (u *Unit) ClearIntegrationTargets(ctx context.Context, projectID, integrati for i, id := range ids { refs[i] = lifecyclelock.AgentRef{ProjectID: projectID, AgentID: id} } - if err = u.LockAgentRefs(ctx, refs, LifecycleAuthority{}); err != nil { + if err = u.LockAgentRefs(ctx, refs, nil); err != nil { return err } return q.ClearExecutionIntegrationTargets( diff --git a/internal/storage/internal/agentexecution/commands_cancel.go b/internal/storage/internal/agentexecution/commands_cancel.go index 91aae5f04..69cb2c3de 100644 --- a/internal/storage/internal/agentexecution/commands_cancel.go +++ b/internal/storage/internal/agentexecution/commands_cancel.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "errors" - "time" "github.com/omnara-ai/omnara/internal/notifications" @@ -45,14 +44,6 @@ func (h *Handle) Cancel(ctx context.Context, input CancelInput) (CanceledExecuti if err != nil { return CanceledExecution{}, err } - if snapshot.View.State == AgentArchived && input.ForceRuntime { - view := snapshot.View - view.State = AgentActive - snapshot.Selection, err = SelectNext(view, time.Time{}) - if err != nil { - return CanceledExecution{}, err - } - } updates := make(map[uuid.UUID]*notifications.AgentInteractionUpdate) turn := uuid.Nil if snapshot.Selection.TurnContinuable { @@ -226,7 +217,6 @@ func (h *Handle) Cancel(ctx context.Context, input CancelInput) (CanceledExecuti } } m.head.StopSequence = result.Event.Sequence - m.head.AnsweredThroughSequence = max(m.head.AnsweredThroughSequence, result.Event.Sequence) m.head.PendingToolOutputID = uuid.Nil m.head.PendingOutputLimitID = uuid.Nil m.head.PendingConfigInputID = uuid.Nil diff --git a/internal/storage/internal/agentexecution/commands_family_integration_test.go b/internal/storage/internal/agentexecution/commands_family_integration_test.go index 523b7d3b2..048271ce5 100644 --- a/internal/storage/internal/agentexecution/commands_family_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_family_integration_test.go @@ -6,10 +6,13 @@ import ( "context" "errors" "testing" + "time" "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/interactionform" "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/omnara-ai/omnara/internal/storage" + "github.com/omnara-ai/omnara/internal/storage/executionstore" "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/stretchr/testify/require" @@ -17,8 +20,8 @@ import ( func TestExecutionCreatedAgentPublicationAndSavepoint(t *testing.T) { f := migratedExecutionFixture(t) - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(t.Context()) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(t.Context(), cell) require.NoError(t, err) defer func() { _ = u.Rollback(context.Background()) }() h, err := u.CreateAgent( @@ -88,8 +91,8 @@ func familyFixture( uuid.New(), ) require.NoError(t, err) - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(t.Context()) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(t.Context(), cell) require.NoError(t, err) t.Cleanup(func() { _ = u.Rollback(context.Background()) }) require.NoError( @@ -100,15 +103,15 @@ func familyFixture( {ProjectID: f.project, AgentID: parent}, {ProjectID: f.project, AgentID: child}, }, - agentexecution.LifecycleAuthority{}, + nil, ), ) p, err := u.Agent( - agentexecution.AgentRoute{CellID: "test", ProjectID: f.project, AgentID: parent, RootAgentID: parent}, + agentexecution.AgentRoute{ProjectID: f.project, AgentID: parent, RootAgentID: parent}, ) require.NoError(t, err) c, err := u.Agent( - agentexecution.AgentRoute{CellID: "test", ProjectID: f.project, AgentID: child, RootAgentID: parent}, + agentexecution.AgentRoute{ProjectID: f.project, AgentID: child, RootAgentID: parent}, ) require.NoError(t, err) return f, u, p, c, lease @@ -315,3 +318,37 @@ func TestExecutionParentReportPreservesInteraction(t *testing.T) { require.Equal(t, agentexecution.WaitToolBatch, snapshot.Selection.Wait) require.NoError(t, u.Commit(t.Context(), "parent report")) } + +func TestSiblingPreparationDoesNotLockParent(t *testing.T) { + f := migratedExecutionFixture(t) + parent := f.create(t, uuid.New(), uuid.Nil) + store := storage.NewStore(f.pool) + var inputs []executionstore.PrepareNormalModelCallInput + for range 8 { + child := f.create(t, uuid.New(), parent) + u, h := f.handle(t, child) + received := receiveCommand(t, h, "queued", "work") + claimed, err := h.Claim(t.Context(), uuid.New(), 90*time.Second) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "claim sibling")) + inputs = append(inputs, + executionstore.PrepareNormalModelCallInput{ProjectID: f.project, + AgentID: child, + RuntimeLockID: claimed.RuntimeID, + OpeningInputIDs: []uuid.UUID{received.ID}}) + } + blocker, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = blocker.Rollback(t.Context()) }() + _, err = blocker.Exec(t.Context(), `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, parent) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + done := make(chan error, len(inputs)) + for _, input := range inputs { + go func() { _, err := store.Execution().PrepareNormalModelCall(ctx, input); done <- err }() + } + for range inputs { + require.NoError(t, <-done) + } +} diff --git a/internal/storage/internal/agentexecution/commands_integration_test.go b/internal/storage/internal/agentexecution/commands_integration_test.go index bb9242a3b..819622f73 100644 --- a/internal/storage/internal/agentexecution/commands_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_integration_test.go @@ -224,7 +224,7 @@ func TestExecutionCommandSavepointAndReplay(t *testing.T) { assertCommandState(t, u, h) return sentinel }), sentinel) - _, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + _, count, _, err := agentexecution.CountOrdinaryLoad(t.Context(), h) require.NoError(t, err) require.Zero(t, count) got := assertCommandState(t, u, h) @@ -777,14 +777,14 @@ func prepareBlockedCandidate(ctx context.Context, a, lease uuid.UUID, decision agentexecution.ModelDecision) error { - cell := agentexecution.NewCell("test", f.pool, nil, nil) - unit, err := cell.Begin(ctx) + cell := agentexecution.NewCell(f.pool, nil, nil) + unit, err := agentexecution.BeginTestUnit(ctx, cell) if err != nil { return err } defer func() { _ = unit.Rollback(context.WithoutCancel(ctx)) }() plan, err := unit.PlanAgents(ctx, []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: a}}, - agentexecution.IngressAuthority{}) + nil) if err != nil { return err } @@ -792,8 +792,7 @@ func prepareBlockedCandidate(ctx context.Context, return err } handle, - err := unit.Agent(agentexecution.AgentRoute{CellID: cell.ID(), - ProjectID: f.project, + err := unit.Agent(agentexecution.AgentRoute{ProjectID: f.project, AgentID: a, RootAgentID: a}) if err != nil { @@ -978,10 +977,41 @@ func TestExecutionPreparationOwnsCachedLineage(t *testing.T) { require.NoError(t, err) selected.Selection.Model.Opening.InputIDs[0] = uuid.New() prepared.Context.Opening.InputIDs[0] = uuid.New() - snapshot, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + snapshot, count, _, err := agentexecution.CountOrdinaryLoad(t.Context(), h) require.NoError(t, err) require.Zero(t, count) require.Equal(t, []uuid.UUID{input.ID}, snapshot.View.NormalContext.Opening.InputIDs) snapshot.View.NormalContext.Opening.InputIDs[0] = uuid.New() assertCommandState(t, u, h) } + +func TestExecutionArchiveHistoryDoesNotDependOnReads(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + receiveCommand(t, h, "queued", "initial") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "admit")) + u, h = f.handle(t, a) + rollback := errors.New("rollback archive") + var histories []string + for _, loaded := range []bool{false, true} { + require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { + if loaded { + _, err = h.LoadExecution(t.Context()) + require.NoError(t, err) + } + _, err = h.Archive(t.Context(), uuid.Nil) + require.NoError(t, err) + var history string + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT coalesce(jsonb_agg(jsonb_build_object('event',to_jsonb(e)-ARRAY['id','agent_input_id','idempotency_key','created_at'],'input',to_jsonb(i)-ARRAY['id','queued_at','admitted_at','resolved_at','canceled_at','admitted_event_id']) ORDER BY e.sequence),'[]')::text FROM agent_events e LEFT JOIN agent_inputs i ON i.id=e.agent_input_id WHERE e.agent_id=$1`, + a).Scan(&history)) + histories = append(histories, history) + return rollback + }), rollback) + } + require.Equal(t, histories[0], histories[1]) + require.Contains(t, histories[0], `"input_kind": "control"`) +} diff --git a/internal/storage/internal/agentexecution/commands_lifecycle.go b/internal/storage/internal/agentexecution/commands_lifecycle.go index 5e17eb34f..90d7c339a 100644 --- a/internal/storage/internal/agentexecution/commands_lifecycle.go +++ b/internal/storage/internal/agentexecution/commands_lifecycle.go @@ -206,7 +206,12 @@ func (h *Handle) Archive(ctx context.Context, actorID uuid.UUID) (bool, error) { if identity.State == "archived" { return false, nil } - if err = h.unit.CaptureWebhookEligibility(ctx, h.route.AgentID); err != nil { + if _, + err = h.Cancel(ctx, + CancelInput{ActorID: actorID, + Reason: "agent_archived", + Message: "The model call was stopped because the agent was archived.", + ForceRuntime: true}); err != nil { return false, err } n, err := executiondb.New(). @@ -244,14 +249,6 @@ func (h *Handle) Archive(ctx context.Context, actorID uuid.UUID) (bool, error) { return false, err } } - if _, - err = h.Cancel(ctx, - CancelInput{ActorID: actorID, - Reason: "agent_archived", - Message: "The model call was stopped because the agent was archived.", - ForceRuntime: true}); err != nil { - return false, err - } if _, err = executiondb.New().CancelExecutionInputs(ctx, h.unit.DB(), diff --git a/internal/storage/internal/agentexecution/commands_models.go b/internal/storage/internal/agentexecution/commands_models.go index d3f07bf74..7150c0d9f 100644 --- a/internal/storage/internal/agentexecution/commands_models.go +++ b/internal/storage/internal/agentexecution/commands_models.go @@ -34,7 +34,7 @@ func (h *Handle) FenceRuntime(ctx context.Context, id uuid.UUID) error { if err := h.unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: h.route.ProjectID, AgentID: h.route.AgentID}}, - RuntimeAuthority{AgentID: h.route.AgentID, RuntimeLockID: id}); err != nil { + &RuntimeAuthority{AgentID: h.route.AgentID, RuntimeLockID: id}); err != nil { return err } _, err := executiondb.New().FenceExecutionRuntime(ctx, h.unit.DB(), executiondb.FenceExecutionRuntimeParams{ diff --git a/internal/storage/internal/agentexecution/commands_outputs.go b/internal/storage/internal/agentexecution/commands_outputs.go index cc6bf17eb..9ddb9fa86 100644 --- a/internal/storage/internal/agentexecution/commands_outputs.go +++ b/internal/storage/internal/agentexecution/commands_outputs.go @@ -334,7 +334,6 @@ func applyAcceptedOutput( event ExecutionEvent, tools, limit bool, ) error { - m.head.AnsweredThroughSequence = max(m.head.AnsweredThroughSequence, event.Sequence) if snapshot.View.NormalContext != nil && snapshot.View.NormalContext.ID == source.ID { snapshot.View.NormalContext = &source } diff --git a/internal/storage/internal/agentexecution/commands_runtime.go b/internal/storage/internal/agentexecution/commands_runtime.go index 98494b1cd..b5715f651 100644 --- a/internal/storage/internal/agentexecution/commands_runtime.go +++ b/internal/storage/internal/agentexecution/commands_runtime.go @@ -161,22 +161,6 @@ func (h *Handle) Advance( }) } -func (h *Handle) AcceptAndAdvance( - ctx context.Context, - input AcceptOutputInput, - allowModelWork, prepareModel bool, -) (AcceptedOutput, OwnedWork, error) { - output, err := h.AcceptOutput(ctx, input) - if err != nil { - return AcceptedOutput{}, OwnedWork{}, err - } - if !output.Created { - return output, OwnedWork{Released: true}, h.Release(ctx, input.RuntimeLockID) - } - work, err := h.Advance(ctx, input.RuntimeLockID, allowModelWork, prepareModel) - return output, work, err -} - func (h *Handle) lockRuntime( ctx context.Context, id uuid.UUID, @@ -252,7 +236,7 @@ type RuntimeRenewal struct { func (u *Unit) RenewRuntime( ctx context.Context, route AgentRoute, id uuid.UUID, leaseDuration time.Duration, ) (RuntimeRenewal, error) { - if route.CellID != u.cell.id || route.ProjectID == uuid.Nil || route.AgentID == uuid.Nil || + if route.ProjectID == uuid.Nil || route.AgentID == uuid.Nil || id == uuid.Nil { return RuntimeRenewal{}, storeerr.ErrRuntimeLockInactive } diff --git a/internal/storage/internal/agentexecution/commands_runtime_integration_test.go b/internal/storage/internal/agentexecution/commands_runtime_integration_test.go index e27b056aa..339e45193 100644 --- a/internal/storage/internal/agentexecution/commands_runtime_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_runtime_integration_test.go @@ -120,14 +120,16 @@ func TestExecutionClaimRenewAcceptAdvance(t *testing.T) { ContextID: prepared.Context.ID, Response: commandResponse(modelenvelope.StopReasonMaxTokens), } - accepted, work, err := h.AcceptAndAdvance(t.Context(), input, true, true) + accepted, err := h.AcceptOutput(t.Context(), input) + require.NoError(t, err) + work, err := h.Advance(t.Context(), input.RuntimeLockID, true, true) require.NoError(t, err) require.True(t, accepted.Created) require.NotEqual(t, prepared.Context.ID, work.Model.Context.ID) - replay, work, err := h.AcceptAndAdvance(t.Context(), input, true, true) + replay, err := h.AcceptOutput(t.Context(), input) require.NoError(t, err) require.False(t, replay.Created) - require.Equal(t, uuid.Nil, work.Model.Context.ID) + require.NoError(t, h.Release(t.Context(), input.RuntimeLockID)) assertCommandState(t, u, h) require.NoError(t, u.Commit(t.Context(), "claim renew advance")) } @@ -192,10 +194,10 @@ func TestExecutionRenewalSamplesTimeAfterLock(t *testing.T) { defer func() { _ = blocker.Rollback(context.Background()) }() _, err = blocker.Exec(t.Context(), `SELECT id FROM agent_runtime_locks WHERE id=$1 FOR UPDATE`, lease) require.NoError(t, err) - cell := agentexecution.NewCell("test", f.pool, nil, nil) + cell := agentexecution.NewCell(f.pool, nil, nil) result := make(chan error, 1) go func() { - u, err := cell.Begin(t.Context()) + u, err := agentexecution.BeginTestUnit(t.Context(), cell) if err != nil { result <- err return @@ -203,8 +205,7 @@ func TestExecutionRenewalSamplesTimeAfterLock(t *testing.T) { defer func() { _ = u.Rollback(context.Background()) }() _, err = u.RenewRuntime(t.Context(), - agentexecution.AgentRoute{CellID: "test", - ProjectID: f.project, + agentexecution.AgentRoute{ProjectID: f.project, AgentID: a, RootAgentID: a}, lease, diff --git a/internal/storage/internal/agentexecution/commands_tools.go b/internal/storage/internal/agentexecution/commands_tools.go index b7c2fa4e9..63ef91121 100644 --- a/internal/storage/internal/agentexecution/commands_tools.go +++ b/internal/storage/internal/agentexecution/commands_tools.go @@ -318,7 +318,7 @@ const ( func (h *Handle) DispatchTool(ctx context.Context, ref ToolRef, acceptExisting bool) (ToolDispatch, error) { if err := h.unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: h.route.ProjectID, AgentID: h.route.AgentID}}, - RuntimeAuthority{AgentID: h.route.AgentID, RuntimeLockID: ref.RuntimeLockID}, + &RuntimeAuthority{AgentID: h.route.AgentID, RuntimeLockID: ref.RuntimeLockID}, ); err != nil { return 0, err } diff --git a/internal/storage/internal/agentexecution/facts_batch.go b/internal/storage/internal/agentexecution/facts_batch.go new file mode 100644 index 000000000..94b55a422 --- /dev/null +++ b/internal/storage/internal/agentexecution/facts_batch.go @@ -0,0 +1,121 @@ +package agentexecution + +import ( + "context" + "errors" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" +) + +type factsBatch struct { + executiondb.DBTX + queued pgx.Batch + results pgx.BatchResults + readers []func() + err error +} + +type factsResults struct{ batch *factsBatch } + +func (b *factsBatch) SendBatch(_ context.Context, queries *pgx.Batch) pgx.BatchResults { + b.queued.QueuedQueries = append(b.queued.QueuedQueries, queries.QueuedQueries...) + return factsResults{batch: b} +} +func (r factsResults) Query() (pgx.Rows, error) { + rows, err := r.batch.results.Query() + if err != nil { + return nil, err + } + if err = rows.Err(); err != nil { + rows.Close() + return nil, err + } + return rows, nil +} +func (r factsResults) QueryRow() pgx.Row { return r.batch.results.QueryRow() } +func (r factsResults) Exec() (pgconn.CommandTag, error) { return r.batch.results.Exec() } +func (r factsResults) Close() error { return nil } + +func (b *factsBatch) read(ctx context.Context) error { + if b.queued.Len() == 0 { + return nil + } + b.results = b.DBTX.SendBatch(ctx, &b.queued) + defer func() { _ = b.results.Close() }() + for _, read := range b.readers { + read() + } + return errors.Join(b.err, b.results.Close()) +} + +type executionFacts struct { + scope executiondb.LoadExecutionFactsRow + contexts []executiondb.LoadExecutionContextsRow + configs []executiondb.LoadExecutionPendingConfigRow + checkpoints []executiondb.LoadExecutionPendingCheckpointRow +} + +func (b *factsBatch) queue(ctx context.Context, route AgentRoute, head ExecutionHead) *executionFacts { + facts := &executionFacts{} + q := executiondb.New() + scope := q.LoadExecutionFacts(ctx, + b, + []executiondb.LoadExecutionFactsParams{{AgentID: route.AgentID, + ProjectID: route.ProjectID, + TurnID: nullableID(head.CurrentTurnID), + StopSequence: head.StopSequence, + MaxContextInputSequence: head.MaxContextInputSequence, + ToolOutputID: nullableID(head.PendingToolOutputID)}}) + b.readers = append(b.readers, func() { + scope.QueryRow(func(_ int, row executiondb.LoadExecutionFactsRow, err error) { + facts.scope = row + b.err = errors.Join(b.err, err) + }) + b.err = errors.Join(b.err, scope.Close()) + }) + if head.NormalContextID != uuid.Nil || + head.CompactionContextID != uuid.Nil || + head.PendingToolOutputID != uuid.Nil || + head.PendingOutputLimitID != uuid.Nil { + rows := q.LoadExecutionContexts(ctx, b, []executiondb.LoadExecutionContextsParams{{AgentID: route.AgentID, + NormalContextID: nullableID(head.NormalContextID), CompactionContextID: nullableID(head.CompactionContextID), + ToolOutputID: nullableID(head.PendingToolOutputID), LimitOutputID: nullableID(head.PendingOutputLimitID)}}) + b.readers = append(b.readers, func() { + rows.Query(func(_ int, rows []executiondb.LoadExecutionContextsRow, err error) { + facts.contexts = rows + b.err = errors.Join(b.err, err) + }) + b.err = errors.Join(b.err, rows.Close()) + }) + } + if head.PendingConfigInputID != uuid.Nil { + rows := q.LoadExecutionPendingConfig(ctx, + b, + []executiondb.LoadExecutionPendingConfigParams{{AgentID: route.AgentID, + ID: head.PendingConfigInputID}}) + b.readers = append(b.readers, func() { + rows.Query(func(_ int, rows []executiondb.LoadExecutionPendingConfigRow, err error) { + facts.configs = rows + b.err = errors.Join(b.err, err) + }) + b.err = errors.Join(b.err, rows.Close()) + }) + } + if head.PendingCheckpointID != uuid.Nil { + rows := q.LoadExecutionPendingCheckpoint(ctx, + b, + []executiondb.LoadExecutionPendingCheckpointParams{{AgentID: route.AgentID, + ID: head.PendingCheckpointID}}) + b.readers = append(b.readers, func() { + rows.Query(func(_ int, rows []executiondb.LoadExecutionPendingCheckpointRow, err error) { + facts.checkpoints = rows + b.err = errors.Join(b.err, err) + }) + b.err = errors.Join(b.err, rows.Close()) + }) + } + return facts +} diff --git a/internal/storage/internal/agentexecution/fixture_integration_test.go b/internal/storage/internal/agentexecution/fixture_integration_test.go index 9c2596440..97707bda9 100644 --- a/internal/storage/internal/agentexecution/fixture_integration_test.go +++ b/internal/storage/internal/agentexecution/fixture_integration_test.go @@ -64,15 +64,15 @@ func seedExecutionFixture(t *testing.T, pool *pgxpool.Pool) executionFixture { func (f executionFixture) create(t *testing.T, id, parent uuid.UUID) uuid.UUID { t.Helper() - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(t.Context()) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(t.Context(), cell) require.NoError(t, err) defer func() { _ = u.Rollback(t.Context()) }() key := "" if parent != uuid.Nil { key = id.String() require.NoError(t, u.LockAgentRefs(t.Context(), - []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: parent}}, agentexecution.LifecycleAuthority{})) + []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: parent}}, nil)) } _, err = u.CreateAgent(t.Context(), agentexecution.CreateAgentInput{ID: id, ProjectID: f.project, ConfigID: f.config, ParentID: parent, SubagentKey: key, Name: "Agent"}) diff --git a/internal/storage/internal/agentexecution/handle.go b/internal/storage/internal/agentexecution/handle.go index ce5d41c59..f56e657e0 100644 --- a/internal/storage/internal/agentexecution/handle.go +++ b/internal/storage/internal/agentexecution/handle.go @@ -1,10 +1,9 @@ package agentexecution import ( - "fmt" - "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" ) type Handle struct { @@ -23,7 +22,7 @@ func (u *Unit) Agent(route AgentRoute) (*Handle, error) { } h, found := u.handles[route.AgentID] if !found || h.route != route { - return nil, fmt.Errorf("%w: agent is not locked in this unit", ErrLockPlan) + return nil, u.expandPlan(lifecyclelock.AgentRef{ProjectID: route.ProjectID, AgentID: route.AgentID}) } return h, nil } @@ -44,7 +43,7 @@ func (u *Unit) Handle(projectID, agentID uuid.UUID) (*Handle, error) { } h := u.handles[agentID] if h == nil || h.route.ProjectID != projectID { - return nil, fmt.Errorf("%w: agent is not locked in this unit", ErrLockPlan) + return nil, u.expandPlan(lifecyclelock.AgentRef{ProjectID: projectID, AgentID: agentID}) } return h, nil } diff --git a/internal/storage/internal/agentexecution/internal/executiondb/batch.go b/internal/storage/internal/agentexecution/internal/executiondb/batch.go index 4115d5f8a..067f51e30 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/batch.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/batch.go @@ -19,8 +19,144 @@ var ( ErrBatchAlreadyClosed = errors.New("batch already closed") ) -const loadExecutionFacts = `-- name: LoadExecutionFacts :batchmany -WITH scope AS MATERIALIZED ( +const loadExecutionContexts = `-- name: LoadExecutionContexts :batchmany +WITH outputs AS ( +SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at +FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id +WHERE o.agent_id=$1 AND o.id IN($4::uuid,$5::uuid) +) +SELECT c.id,c.agent_id,c.turn_id,c.operation_kind,c.attempt_number,c.input_event_sequence,c.source_event_sequence_end, +c.state,c.recovery_kind,c.retry_at,c.opening_input_ids,c.opening_event_sequence, +o.id AS output_id,o.stop_reason AS output_stop_reason,o.event_id AS output_event_id, +o.sequence AS output_sequence,o.created_at AS output_created_at +FROM model_call_contexts c LEFT JOIN outputs o ON o.model_call_context_id=c.id +WHERE c.agent_id=$1 AND c.id=ANY( +ARRAY[$2::uuid,$3::uuid] +|| ARRAY(SELECT model_call_context_id FROM outputs)) +` + +type LoadExecutionContextsBatchResults struct { + br pgx.BatchResults + tot int + closed bool +} + +type LoadExecutionContextsParams struct { + AgentID uuid.UUID + NormalContextID *uuid.UUID + CompactionContextID *uuid.UUID + ToolOutputID *uuid.UUID + LimitOutputID *uuid.UUID +} + +type LoadExecutionContextsRow struct { + ID uuid.UUID + AgentID uuid.UUID + TurnID uuid.UUID + OperationKind string + AttemptNumber int32 + InputEventSequence int64 + SourceEventSequenceEnd *int64 + State string + RecoveryKind *string + RetryAt *time.Time + OpeningInputIds []uuid.UUID + OpeningEventSequence int64 + OutputID *uuid.UUID + OutputStopReason *string + OutputEventID *uuid.UUID + OutputSequence *int64 + OutputCreatedAt *time.Time +} + +func (q *Queries) LoadExecutionContexts(ctx context.Context, db DBTX, arg []LoadExecutionContextsParams) *LoadExecutionContextsBatchResults { + batch := &pgx.Batch{} + for _, a := range arg { + vals := []interface{}{ + a.AgentID, + a.NormalContextID, + a.CompactionContextID, + a.ToolOutputID, + a.LimitOutputID, + } + batch.Queue(loadExecutionContexts, vals...) + } + br := db.SendBatch(ctx, batch) + return &LoadExecutionContextsBatchResults{br, len(arg), false} +} + +func (b *LoadExecutionContextsBatchResults) Query(f func(int, []LoadExecutionContextsRow, error)) { + defer b.br.Close() + for t := 0; t < b.tot; t++ { + items := []LoadExecutionContextsRow{} + if b.closed { + if f != nil { + f(t, items, ErrBatchAlreadyClosed) + } + continue + } + err := func() error { + rows, err := b.br.Query() + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var i LoadExecutionContextsRow + if err := rows.Scan( + &i.ID, + &i.AgentID, + &i.TurnID, + &i.OperationKind, + &i.AttemptNumber, + &i.InputEventSequence, + &i.SourceEventSequenceEnd, + &i.State, + &i.RecoveryKind, + &i.RetryAt, + &i.OpeningInputIds, + &i.OpeningEventSequence, + &i.OutputID, + &i.OutputStopReason, + &i.OutputEventID, + &i.OutputSequence, + &i.OutputCreatedAt, + ); err != nil { + return err + } + items = append(items, i) + } + return rows.Err() + }() + if f != nil { + f(t, items, err) + } + } +} + +func (b *LoadExecutionContextsBatchResults) Close() error { + b.closed = true + return b.br.Close() +} + +const loadExecutionFacts = `-- name: LoadExecutionFacts :batchone +WITH batch AS ( +SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2) AS has_tools, +EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2 AND c.state<>'completed') AS incomplete, +EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$1 AND c.model_output_id=$2 +AND (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable FROM (SELECT 1 WHERE $2::uuid IS NOT NULL) present +), tool_batch AS ( +SELECT b.has_tools,b.incomplete,b.runnable,coalesce(r.calls,0)::bigint AS calls,coalesce(r.results,0)::bigint AS results, +coalesce(r.events,0)::bigint AS events,coalesce(r.last_sequence,0)::bigint AS last_sequence, +coalesce(r.ready_at,'epoch'::timestamptz)::timestamptz AS ready_at +FROM batch b LEFT JOIN LATERAL ( +SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, +max(e.sequence)::bigint AS last_sequence,max(r.completed_at)::timestamptz AS ready_at +FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id +LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id +WHERE c.agent_id=$1 AND c.model_output_id=$2 AND NOT b.incomplete +) r ON NOT b.incomplete +), scope AS MATERIALIZED ( SELECT a.id,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now,t.id AS turn_id, coalesce(e.sequence,0)::bigint AS semantic_sequence,coalesce(e.created_at,'epoch'::timestamptz)::timestamptz AS semantic_time, EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' @@ -30,7 +166,7 @@ WITH scope AS MATERIALIZED ( coalesce(o.first_opening_sequence,0)::bigint AS first_opening_sequence, coalesce(o.last_opening_sequence,0)::bigint AS last_opening_sequence, coalesce(o.first_content_sequence,0)::bigint AS first_content_sequence - FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=$1 + FROM agents a LEFT JOIN agent_turns t ON t.agent_id=a.id AND t.id=$3 LEFT JOIN agent_events e ON e.agent_id=a.id AND e.id=t.latest_semantic_event_id LEFT JOIN LATERAL ( SELECT min(x.sequence) AS first_opening_sequence,max(x.sequence) AS last_opening_sequence, @@ -38,29 +174,16 @@ WITH scope AS MATERIALIZED ( FROM agent_events x JOIN agent_inputs i ON i.agent_id=x.agent_id AND i.id=x.agent_input_id WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event ) o ON true - WHERE a.id=$2 AND a.project_id=$3 -), outputs AS MATERIALIZED ( - SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at - FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id - WHERE o.agent_id=$2 AND o.id IN ($4::uuid,$5::uuid) -), batch AS MATERIALIZED ( - SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$2 - AND c.model_output_id=$4) AS has_tools, - EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$2 - AND c.model_output_id=$4 AND c.state<>'completed') AS incomplete, - EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=$2 - AND c.model_output_id=$4 AND - (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable - WHERE $4::uuid IS NOT NULL -), boundary AS MATERIALIZED ( - SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$2 + WHERE a.id=$1 AND a.project_id=$4 +) , boundary AS MATERIALIZED ( + SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$1 AND e.event_kind='model_output' AND e.sequence<=(SELECT last_opening_sequence FROM scope) - ORDER BY e.sequence DESC LIMIT 1),0),$6::bigint)::bigint AS sequence - FROM scope WHERE scope.state='active' AND $7::bigintb.sequence AND e.sequence<=(SELECT last_opening_sequence FROM scope) JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id @@ -68,208 +191,23 @@ WITH scope AS MATERIALIZED ( SELECT id,sequence,created_at FROM unanswered UNION ALL SELECT i.id,e.sequence,e.created_at FROM boundary b - JOIN agent_events e ON e.agent_id=$2 AND e.turn_id=$1 AND e.is_opening_event + JOIN agent_events e ON e.agent_id=$1 AND e.turn_id=$3 AND e.is_opening_event AND e.sequence<=(SELECT last_opening_sequence FROM scope) JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id AND i.input_kind='content' AND i.state='resolved' AND i.admitted_event_id=e.id WHERE NOT EXISTS(SELECT 1 FROM unanswered) -), facts AS ( -SELECT 'context'::text AS kind, - c.id AS id, - c.turn_id AS turn_id, - '00000000-0000-0000-0000-000000000000'::uuid AS context_id, - '00000000-0000-0000-0000-000000000000'::uuid AS event_id, - 0::bigint AS sequence, - 'epoch'::timestamptz AS event_time, - c.operation_kind AS operation, - c.attempt_number AS attempt, - c.input_event_sequence AS watermark, - coalesce(c.source_event_sequence_end,0)::bigint AS source_end, - c.state AS state, - coalesce(c.recovery_kind,'')::text AS recovery, - c.retry_at AS retry_at, - c.opening_input_ids AS opening_ids, - c.opening_event_sequence AS opening_sequence, - false AS has_tools, - false AS incomplete, - false AS runnable, - 0::bigint AS calls, - 0::bigint AS results, - 0::bigint AS events, - ''::text AS stop_reason, - '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, - 'epoch'::timestamptz AS database_now, - false AS steering,false AS queued, - 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM model_call_contexts c WHERE c.agent_id=$2 AND c.id=ANY( - ARRAY[$8::uuid,$9::uuid] - || ARRAY(SELECT model_call_context_id FROM outputs)) -UNION ALL -SELECT 'output'::text AS kind, - o.id AS id, - '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, - o.model_call_context_id AS context_id, - o.event_id AS event_id, - o.sequence AS sequence, - o.created_at AS event_time, - ''::text AS operation, - 0::integer AS attempt, - 0::bigint AS watermark, - 0::bigint AS source_end, - ''::text AS state, - ''::text AS recovery, - NULL::timestamptz AS retry_at, - '{}'::uuid[] AS opening_ids, - 0::bigint AS opening_sequence, - false AS has_tools, - false AS incomplete, - false AS runnable, - 0::bigint AS calls, - 0::bigint AS results, - 0::bigint AS events, - o.stop_reason AS stop_reason, - '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, - 'epoch'::timestamptz AS database_now, - false AS steering,false AS queued, - 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM outputs o -UNION ALL -SELECT 'batch'::text AS kind, - '00000000-0000-0000-0000-000000000000'::uuid AS id, - '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, - '00000000-0000-0000-0000-000000000000'::uuid AS context_id, - '00000000-0000-0000-0000-000000000000'::uuid AS event_id, - coalesce(r.last_sequence,0)::bigint AS sequence, - coalesce(r.ready_at,'epoch'::timestamptz)::timestamptz AS event_time, - ''::text AS operation, - 0::integer AS attempt, - 0::bigint AS watermark, - 0::bigint AS source_end, - ''::text AS state, - ''::text AS recovery, - NULL::timestamptz AS retry_at, - '{}'::uuid[] AS opening_ids, - 0::bigint AS opening_sequence, - b.has_tools AS has_tools, - b.incomplete AS incomplete, - b.runnable AS runnable, - coalesce(r.calls,0)::bigint AS calls, - coalesce(r.results,0)::bigint AS results, - coalesce(r.events,0)::bigint AS events, - ''::text AS stop_reason, - '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, - 'epoch'::timestamptz AS database_now, - false AS steering,false AS queued, - 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM batch b LEFT JOIN LATERAL ( - SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, - max(e.sequence)::bigint AS last_sequence,max(r.completed_at)::timestamptz AS ready_at - FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id - LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id - WHERE c.agent_id=$2 AND c.model_output_id=$4 AND NOT b.incomplete -) r ON NOT b.incomplete -UNION ALL -SELECT 'config'::text AS kind, - i.id AS id, - e.turn_id AS turn_id, - '00000000-0000-0000-0000-000000000000'::uuid AS context_id, - '00000000-0000-0000-0000-000000000000'::uuid AS event_id, - e.sequence AS sequence, - e.created_at AS event_time, - ''::text AS operation, - 0::integer AS attempt, - 0::bigint AS watermark, - 0::bigint AS source_end, - ''::text AS state, - ''::text AS recovery, - NULL::timestamptz AS retry_at, - i.opening_input_ids AS opening_ids, - coalesce(i.opening_event_sequence,0)::bigint AS opening_sequence, - false AS has_tools, - false AS incomplete, - false AS runnable, - 0::bigint AS calls, - 0::bigint AS results, - 0::bigint AS events, - ''::text AS stop_reason, - '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, - 'epoch'::timestamptz AS database_now, - false AS steering,false AS queued, - 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id -WHERE i.agent_id=$2 AND i.id=$10 AND i.input_kind='config_change' AND i.state='resolved' -UNION ALL -SELECT 'checkpoint'::text AS kind, - c.id AS id, - e.turn_id AS turn_id, - '00000000-0000-0000-0000-000000000000'::uuid AS context_id, - '00000000-0000-0000-0000-000000000000'::uuid AS event_id, - e.sequence AS sequence, - c.created_at AS event_time, - ''::text AS operation, - 0::integer AS attempt, - 0::bigint AS watermark, - 0::bigint AS source_end, - ''::text AS state, - ''::text AS recovery, - NULL::timestamptz AS retry_at, - c.opening_input_ids AS opening_ids, - coalesce(c.opening_event_sequence,0)::bigint AS opening_sequence, - false AS has_tools, - false AS incomplete, - false AS runnable, - 0::bigint AS calls, - 0::bigint AS results, - 0::bigint AS events, - ''::text AS stop_reason, - '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, - 'epoch'::timestamptz AS database_now, - false AS steering,false AS queued, - 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id -WHERE c.agent_id=$2 AND c.id=$11 -UNION ALL -SELECT 'opening'::text AS kind, - o.id AS id, - '00000000-0000-0000-0000-000000000000'::uuid AS turn_id, - '00000000-0000-0000-0000-000000000000'::uuid AS context_id, - '00000000-0000-0000-0000-000000000000'::uuid AS event_id, - o.sequence AS sequence, - o.created_at AS event_time, - ''::text AS operation, - 0::integer AS attempt, - 0::bigint AS watermark, - 0::bigint AS source_end, - ''::text AS state, - ''::text AS recovery, - NULL::timestamptz AS retry_at, - '{}'::uuid[] AS opening_ids, - 0::bigint AS opening_sequence, - false AS has_tools, - false AS incomplete, - false AS runnable, - 0::bigint AS calls, - 0::bigint AS results, - 0::bigint AS events, - ''::text AS stop_reason, - '00000000-0000-0000-0000-000000000000'::uuid AS root_agent_id, - 'epoch'::timestamptz AS database_now, - false AS steering,false AS queued, - 0::bigint AS first_opening_sequence,0::bigint AS last_opening_sequence,0::bigint AS first_content_sequence -FROM opening o -UNION ALL -SELECT 'scope'::text AS kind,s.id,coalesce(s.turn_id,'00000000-0000-0000-0000-000000000000'::uuid), - '00000000-0000-0000-0000-000000000000'::uuid,'00000000-0000-0000-0000-000000000000'::uuid, - s.semantic_sequence,s.semantic_time,''::text,0::integer,0::bigint,0::bigint,s.state,''::text,NULL::timestamptz, - '{}'::uuid[],0::bigint,false,false,false,0::bigint,0::bigint,0::bigint,''::text, - s.root_agent_id,s.database_now,s.steering,s.queued,s.first_opening_sequence,s.last_opening_sequence,s.first_content_sequence -FROM scope s - ) -SELECT f.kind,f.id,f.turn_id,f.context_id,f.event_id,f.sequence,f.event_time,f.operation,f.attempt, - f.watermark,f.source_end,f.state,f.recovery,f.retry_at,f.opening_ids,f.opening_sequence, - f.has_tools,f.incomplete,f.runnable,f.calls,f.results,f.events,f.stop_reason,f.root_agent_id, - f.database_now,f.steering,f.queued,f.first_opening_sequence,f.last_opening_sequence,f.first_content_sequence -FROM facts f ORDER BY f.kind,f.sequence +) +SELECT s.id AS agent_id,s.root_agent_id,s.state,s.database_now,s.turn_id, +s.semantic_sequence,s.semantic_time,s.steering,s.queued,s.first_opening_sequence,s.last_opening_sequence,s.first_content_sequence, +coalesce(o.ids,'{}'::uuid[])::uuid[] AS opening_ids,coalesce(o.first_sequence,0)::bigint AS opening_sequence, +coalesce(o.ready_at,'epoch'::timestamptz)::timestamptz AS opening_ready_at, +coalesce(b.has_tools,false)::boolean AS batch_has_tools, coalesce(b.incomplete,false)::boolean AS batch_incomplete, +coalesce(b.runnable,false)::boolean AS batch_runnable,coalesce(b.calls,0)::bigint AS batch_calls, +coalesce(b.results,0)::bigint AS batch_results,coalesce(b.events,0)::bigint AS batch_events, +coalesce(b.last_sequence,0)::bigint AS batch_last_sequence,coalesce(b.ready_at,'epoch'::timestamptz)::timestamptz AS batch_ready_at +FROM scope s LEFT JOIN LATERAL ( +SELECT array_agg(id ORDER BY sequence) AS ids,min(sequence) AS first_sequence,min(created_at) AS ready_at FROM opening +) o ON true LEFT JOIN tool_batch b ON true ` type LoadExecutionFactsBatchResults struct { @@ -279,67 +217,50 @@ type LoadExecutionFactsBatchResults struct { } type LoadExecutionFactsParams struct { - TurnID *uuid.UUID AgentID uuid.UUID - ProjectID uuid.UUID ToolOutputID *uuid.UUID - LimitOutputID *uuid.UUID + TurnID *uuid.UUID + ProjectID uuid.UUID StopSequence int64 MaxContextInputSequence int64 - NormalContextID *uuid.UUID - CompactionContextID *uuid.UUID - ConfigID *uuid.UUID - CheckpointID *uuid.UUID } type LoadExecutionFactsRow struct { - Kind string - ID uuid.UUID - TurnID uuid.UUID - ContextID uuid.UUID - EventID uuid.UUID - Sequence int64 - EventTime time.Time - Operation string - Attempt int32 - Watermark int64 - SourceEnd int64 - State string - Recovery string - RetryAt *time.Time - OpeningIds []uuid.UUID - OpeningSequence int64 - HasTools bool - Incomplete bool - Runnable bool - Calls int64 - Results int64 - Events int64 - StopReason string + AgentID uuid.UUID RootAgentID uuid.UUID + State string DatabaseNow time.Time + TurnID *uuid.UUID + SemanticSequence int64 + SemanticTime time.Time Steering bool Queued bool FirstOpeningSequence int64 LastOpeningSequence int64 FirstContentSequence int64 + OpeningIds []uuid.UUID + OpeningSequence int64 + OpeningReadyAt time.Time + BatchHasTools bool + BatchIncomplete bool + BatchRunnable bool + BatchCalls int64 + BatchResults int64 + BatchEvents int64 + BatchLastSequence int64 + BatchReadyAt time.Time } func (q *Queries) LoadExecutionFacts(ctx context.Context, db DBTX, arg []LoadExecutionFactsParams) *LoadExecutionFactsBatchResults { batch := &pgx.Batch{} for _, a := range arg { vals := []interface{}{ - a.TurnID, a.AgentID, - a.ProjectID, a.ToolOutputID, - a.LimitOutputID, + a.TurnID, + a.ProjectID, a.StopSequence, a.MaxContextInputSequence, - a.NormalContextID, - a.CompactionContextID, - a.ConfigID, - a.CheckpointID, } batch.Queue(loadExecutionFacts, vals...) } @@ -347,10 +268,96 @@ func (q *Queries) LoadExecutionFacts(ctx context.Context, db DBTX, arg []LoadExe return &LoadExecutionFactsBatchResults{br, len(arg), false} } -func (b *LoadExecutionFactsBatchResults) Query(f func(int, []LoadExecutionFactsRow, error)) { +func (b *LoadExecutionFactsBatchResults) QueryRow(f func(int, LoadExecutionFactsRow, error)) { defer b.br.Close() for t := 0; t < b.tot; t++ { - items := []LoadExecutionFactsRow{} + var i LoadExecutionFactsRow + if b.closed { + if f != nil { + f(t, i, ErrBatchAlreadyClosed) + } + continue + } + row := b.br.QueryRow() + err := row.Scan( + &i.AgentID, + &i.RootAgentID, + &i.State, + &i.DatabaseNow, + &i.TurnID, + &i.SemanticSequence, + &i.SemanticTime, + &i.Steering, + &i.Queued, + &i.FirstOpeningSequence, + &i.LastOpeningSequence, + &i.FirstContentSequence, + &i.OpeningIds, + &i.OpeningSequence, + &i.OpeningReadyAt, + &i.BatchHasTools, + &i.BatchIncomplete, + &i.BatchRunnable, + &i.BatchCalls, + &i.BatchResults, + &i.BatchEvents, + &i.BatchLastSequence, + &i.BatchReadyAt, + ) + if f != nil { + f(t, i, err) + } + } +} + +func (b *LoadExecutionFactsBatchResults) Close() error { + b.closed = true + return b.br.Close() +} + +const loadExecutionPendingCheckpoint = `-- name: LoadExecutionPendingCheckpoint :batchmany +SELECT c.id,e.turn_id,e.sequence,c.created_at,c.opening_input_ids,coalesce(c.opening_event_sequence,0)::bigint AS opening_event_sequence +FROM context_checkpoints c JOIN agent_events e ON e.agent_id=c.agent_id AND e.context_checkpoint_id=c.id +WHERE c.agent_id=$1 AND c.id=$2 +` + +type LoadExecutionPendingCheckpointBatchResults struct { + br pgx.BatchResults + tot int + closed bool +} + +type LoadExecutionPendingCheckpointParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +type LoadExecutionPendingCheckpointRow struct { + ID uuid.UUID + TurnID uuid.UUID + Sequence int64 + CreatedAt time.Time + OpeningInputIds []uuid.UUID + OpeningEventSequence int64 +} + +func (q *Queries) LoadExecutionPendingCheckpoint(ctx context.Context, db DBTX, arg []LoadExecutionPendingCheckpointParams) *LoadExecutionPendingCheckpointBatchResults { + batch := &pgx.Batch{} + for _, a := range arg { + vals := []interface{}{ + a.AgentID, + a.ID, + } + batch.Queue(loadExecutionPendingCheckpoint, vals...) + } + br := db.SendBatch(ctx, batch) + return &LoadExecutionPendingCheckpointBatchResults{br, len(arg), false} +} + +func (b *LoadExecutionPendingCheckpointBatchResults) Query(f func(int, []LoadExecutionPendingCheckpointRow, error)) { + defer b.br.Close() + for t := 0; t < b.tot; t++ { + items := []LoadExecutionPendingCheckpointRow{} if b.closed { if f != nil { f(t, items, ErrBatchAlreadyClosed) @@ -364,38 +371,14 @@ func (b *LoadExecutionFactsBatchResults) Query(f func(int, []LoadExecutionFactsR } defer rows.Close() for rows.Next() { - var i LoadExecutionFactsRow + var i LoadExecutionPendingCheckpointRow if err := rows.Scan( - &i.Kind, &i.ID, &i.TurnID, - &i.ContextID, - &i.EventID, &i.Sequence, - &i.EventTime, - &i.Operation, - &i.Attempt, - &i.Watermark, - &i.SourceEnd, - &i.State, - &i.Recovery, - &i.RetryAt, - &i.OpeningIds, - &i.OpeningSequence, - &i.HasTools, - &i.Incomplete, - &i.Runnable, - &i.Calls, - &i.Results, - &i.Events, - &i.StopReason, - &i.RootAgentID, - &i.DatabaseNow, - &i.Steering, - &i.Queued, - &i.FirstOpeningSequence, - &i.LastOpeningSequence, - &i.FirstContentSequence, + &i.CreatedAt, + &i.OpeningInputIds, + &i.OpeningEventSequence, ); err != nil { return err } @@ -409,23 +392,104 @@ func (b *LoadExecutionFactsBatchResults) Query(f func(int, []LoadExecutionFactsR } } -func (b *LoadExecutionFactsBatchResults) Close() error { +func (b *LoadExecutionPendingCheckpointBatchResults) Close() error { + b.closed = true + return b.br.Close() +} + +const loadExecutionPendingConfig = `-- name: LoadExecutionPendingConfig :batchmany +SELECT i.id,e.turn_id,e.sequence,e.created_at,i.opening_input_ids,coalesce(i.opening_event_sequence,0)::bigint AS opening_event_sequence +FROM agent_inputs i JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id +WHERE i.agent_id=$1 AND i.id=$2 AND i.input_kind='config_change' AND i.state='resolved' +` + +type LoadExecutionPendingConfigBatchResults struct { + br pgx.BatchResults + tot int + closed bool +} + +type LoadExecutionPendingConfigParams struct { + AgentID uuid.UUID + ID uuid.UUID +} + +type LoadExecutionPendingConfigRow struct { + ID uuid.UUID + TurnID uuid.UUID + Sequence int64 + CreatedAt time.Time + OpeningInputIds []uuid.UUID + OpeningEventSequence int64 +} + +func (q *Queries) LoadExecutionPendingConfig(ctx context.Context, db DBTX, arg []LoadExecutionPendingConfigParams) *LoadExecutionPendingConfigBatchResults { + batch := &pgx.Batch{} + for _, a := range arg { + vals := []interface{}{ + a.AgentID, + a.ID, + } + batch.Queue(loadExecutionPendingConfig, vals...) + } + br := db.SendBatch(ctx, batch) + return &LoadExecutionPendingConfigBatchResults{br, len(arg), false} +} + +func (b *LoadExecutionPendingConfigBatchResults) Query(f func(int, []LoadExecutionPendingConfigRow, error)) { + defer b.br.Close() + for t := 0; t < b.tot; t++ { + items := []LoadExecutionPendingConfigRow{} + if b.closed { + if f != nil { + f(t, items, ErrBatchAlreadyClosed) + } + continue + } + err := func() error { + rows, err := b.br.Query() + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var i LoadExecutionPendingConfigRow + if err := rows.Scan( + &i.ID, + &i.TurnID, + &i.Sequence, + &i.CreatedAt, + &i.OpeningInputIds, + &i.OpeningEventSequence, + ); err != nil { + return err + } + items = append(items, i) + } + return rows.Err() + }() + if f != nil { + f(t, items, err) + } + } +} + +func (b *LoadExecutionPendingConfigBatchResults) Close() error { b.closed = true return b.br.Close() } const publishExecution = `-- name: PublishExecution :batchone WITH scope AS MATERIALIZED ( - SELECT a.id AS agent_id,$2::timestamptz AS logical_ready_at - FROM agents a WHERE a.id=$3 AND ($4::boolean OR + SELECT a.id AS agent_id,$1::timestamptz AS logical_ready_at + FROM agents a WHERE a.id=$2 AND ($3::boolean OR EXISTS(SELECT 1 FROM agent_execution_state h WHERE h.agent_id=a.id)) ), updated AS ( -INSERT INTO agent_execution_state AS h(agent_id,current_turn_id,stop_sequence,answered_through_sequence,max_normal_input_sequence,max_context_input_sequence,normal_context_id,compaction_context_id,pending_tool_output_id,pending_output_limit_id,pending_config_input_id,pending_checkpoint_id,turn_continuable,incomplete_tools,logical_ready_at) -SELECT scope.agent_id,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,scope.logical_ready_at FROM scope +INSERT INTO agent_execution_state AS h(agent_id,current_turn_id,stop_sequence,max_normal_input_sequence,max_context_input_sequence,normal_context_id,compaction_context_id,pending_tool_output_id,pending_output_limit_id,pending_config_input_id,pending_checkpoint_id,turn_continuable,incomplete_tools,logical_ready_at) +SELECT scope.agent_id,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,scope.logical_ready_at FROM scope ON CONFLICT(agent_id) DO UPDATE SET current_turn_id=excluded.current_turn_id, stop_sequence=excluded.stop_sequence, -answered_through_sequence=excluded.answered_through_sequence, max_normal_input_sequence=excluded.max_normal_input_sequence, max_context_input_sequence=excluded.max_context_input_sequence, normal_context_id=excluded.normal_context_id, @@ -437,14 +501,14 @@ pending_checkpoint_id=excluded.pending_checkpoint_id, turn_continuable=excluded.turn_continuable, incomplete_tools=excluded.incomplete_tools, logical_ready_at=excluded.logical_ready_at -WHERE (h.current_turn_id,h.stop_sequence,h.answered_through_sequence,h.max_normal_input_sequence,h.max_context_input_sequence,h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id,h.pending_config_input_id,h.pending_checkpoint_id,h.turn_continuable,h.incomplete_tools,h.logical_ready_at) IS DISTINCT FROM -(excluded.current_turn_id,excluded.stop_sequence,excluded.answered_through_sequence,excluded.max_normal_input_sequence,excluded.max_context_input_sequence,excluded.normal_context_id,excluded.compaction_context_id,excluded.pending_tool_output_id,excluded.pending_output_limit_id,excluded.pending_config_input_id,excluded.pending_checkpoint_id,excluded.turn_continuable,excluded.incomplete_tools,excluded.logical_ready_at) +WHERE (h.current_turn_id,h.stop_sequence,h.max_normal_input_sequence,h.max_context_input_sequence,h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id,h.pending_config_input_id,h.pending_checkpoint_id,h.turn_continuable,h.incomplete_tools,h.logical_ready_at) IS DISTINCT FROM +(excluded.current_turn_id,excluded.stop_sequence,excluded.max_normal_input_sequence,excluded.max_context_input_sequence,excluded.normal_context_id,excluded.compaction_context_id,excluded.pending_tool_output_id,excluded.pending_output_limit_id,excluded.pending_config_input_id,excluded.pending_checkpoint_id,excluded.turn_continuable,excluded.incomplete_tools,excluded.logical_ready_at) RETURNING h.agent_id), desired AS MATERIALIZED ( SELECT a.id,h.logical_ready_at AS ready_at FROM scope h JOIN agents a ON a.id=h.agent_id - WHERE a.id=$3 AND a.state='active' AND h.logical_ready_at IS NOT NULL + WHERE a.id=$2 AND a.state='active' AND h.logical_ready_at IS NOT NULL AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks r WHERE r.agent_id=a.id) ), removed AS ( - DELETE FROM agent_wakeups WHERE agent_id=$3 AND NOT EXISTS(SELECT 1 FROM desired) + DELETE FROM agent_wakeups WHERE agent_id=$2 AND NOT EXISTS(SELECT 1 FROM desired) ), published AS ( INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) SELECT id,ready_at,statement_timestamp() FROM desired @@ -465,7 +529,6 @@ FROM scope LEFT JOIN LATERAL ( JOIN projects p ON p.id=a.project_id JOIN orgs o ON o.id=p.org_id WHERE a.id=scope.agent_id AND p.deleted_at IS NULL AND o.deleted_at IS NULL AND coalesce(c.compiled_definition->'event_webhook'->>'url','')<>'' - AND $1::boolean ) w ON true ` @@ -476,13 +539,11 @@ type PublishExecutionBatchResults struct { } type PublishExecutionParams struct { - CaptureWebhook bool LogicalReadyAt *time.Time AgentID uuid.UUID InsertHead bool CurrentTurnID *uuid.UUID StopSequence int64 - AnsweredThroughSequence int64 MaxNormalInputSequence int64 MaxContextInputSequence int64 NormalContextID *uuid.UUID @@ -505,13 +566,11 @@ func (q *Queries) PublishExecution(ctx context.Context, db DBTX, arg []PublishEx batch := &pgx.Batch{} for _, a := range arg { vals := []interface{}{ - a.CaptureWebhook, a.LogicalReadyAt, a.AgentID, a.InsertHead, a.CurrentTurnID, a.StopSequence, - a.AnsweredThroughSequence, a.MaxNormalInputSequence, a.MaxContextInputSequence, a.NormalContextID, diff --git a/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go index 13813c9b4..c7c2dc36d 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/load.sql.go @@ -26,7 +26,6 @@ func (q *Queries) ExecutionDatabaseTime(ctx context.Context, db DBTX) (time.Time const loadExecutionBase = `-- name: LoadExecutionBase :one SELECT coalesce(h.agent_id,a.id)::uuid AS agent_id,h.current_turn_id, coalesce(h.stop_sequence,0)::bigint AS stop_sequence, - coalesce(h.answered_through_sequence,0)::bigint AS answered_through_sequence, coalesce(h.max_normal_input_sequence,0)::bigint AS max_normal_input_sequence, coalesce(h.max_context_input_sequence,0)::bigint AS max_context_input_sequence, h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id, @@ -46,7 +45,6 @@ type LoadExecutionBaseRow struct { AgentID uuid.UUID CurrentTurnID *uuid.UUID StopSequence int64 - AnsweredThroughSequence int64 MaxNormalInputSequence int64 MaxContextInputSequence int64 NormalContextID *uuid.UUID @@ -71,7 +69,6 @@ func (q *Queries) LoadExecutionBase(ctx context.Context, db DBTX, arg LoadExecut &i.AgentID, &i.CurrentTurnID, &i.StopSequence, - &i.AnsweredThroughSequence, &i.MaxNormalInputSequence, &i.MaxContextInputSequence, &i.NormalContextID, diff --git a/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go index 3ce9c759e..123f71ff6 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/state.sql.go @@ -16,8 +16,6 @@ SELECT coalesce((SELECT max(e.sequence) FROM agent_events e JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id WHERE i.project_id=$1 AND i.agent_id=$2 AND i.input_kind='control' AND i.state='resolved' AND i.control_type='cancel_current'),0)::bigint AS stop_sequence, - coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$2 AND e.event_kind='model_output' - ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS output_sequence, coalesce((SELECT t.id FROM agent_turns t WHERE t.agent_id=$2 ORDER BY t.turn_sequence DESC LIMIT 1), '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS turn_id ` @@ -28,15 +26,14 @@ type LoadExecutionBoundariesParams struct { } type LoadExecutionBoundariesRow struct { - StopSequence int64 - OutputSequence int64 - TurnID uuid.UUID + StopSequence int64 + TurnID uuid.UUID } func (q *Queries) LoadExecutionBoundaries(ctx context.Context, db DBTX, arg LoadExecutionBoundariesParams) (LoadExecutionBoundariesRow, error) { row := db.QueryRow(ctx, loadExecutionBoundaries, arg.ProjectID, arg.AgentID) var i LoadExecutionBoundariesRow - err := row.Scan(&i.StopSequence, &i.OutputSequence, &i.TurnID) + err := row.Scan(&i.StopSequence, &i.TurnID) return i, err } diff --git a/internal/storage/internal/agentexecution/linearity_integration_test.go b/internal/storage/internal/agentexecution/linearity_integration_test.go index a3a817435..6ba0b8908 100644 --- a/internal/storage/internal/agentexecution/linearity_integration_test.go +++ b/internal/storage/internal/agentexecution/linearity_integration_test.go @@ -21,8 +21,8 @@ func TestExecutionBulkAndLoopLinearity(t *testing.T) { var baseline time.Duration for _, n := range []int{1000, 10000} { t.Run(fmt.Sprint(n), func(t *testing.T) { - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(ctx) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(ctx, cell) require.NoError(t, err) defer func() { _ = u.Rollback(ctx) }() started := time.Now() @@ -100,19 +100,18 @@ func TestExecutionReconstructionTenThousandAgents(t *testing.T) { } require.NoError(t, rows.Err()) rows.Close() - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(ctx) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(ctx, cell) require.NoError(t, err) defer func() { _ = u.Rollback(ctx) }() started = time.Now() - plan, err := u.PlanAgents(ctx, refs, agentexecution.RepairAuthority{}) + plan, err := u.PlanAgents(ctx, refs, nil) require.NoError(t, err) require.NoError(t, u.LockAgents(ctx, plan)) var firstThousand time.Duration for i, ref := range refs { h, handleErr := u.Agent( agentexecution.AgentRoute{ - CellID: cell.ID(), ProjectID: f.project, AgentID: ref.AgentID, RootAgentID: ref.AgentID, diff --git a/internal/storage/internal/agentexecution/load.go b/internal/storage/internal/agentexecution/load.go index 6d435b6ed..43f39ac28 100644 --- a/internal/storage/internal/agentexecution/load.go +++ b/internal/storage/internal/agentexecution/load.go @@ -2,9 +2,7 @@ package agentexecution import ( "context" - "errors" "fmt" - "time" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -40,132 +38,98 @@ func loadExecution( return loadExecutionFacts(ctx, db, route, executionHead(base)) } -func executionFactsParams(route AgentRoute, head ExecutionHead) executiondb.LoadExecutionFactsParams { - return executiondb.LoadExecutionFactsParams{AgentID: route.AgentID, ProjectID: route.ProjectID, - TurnID: nullableID( - head.CurrentTurnID, - ), StopSequence: head.StopSequence, MaxContextInputSequence: head.MaxContextInputSequence, - NormalContextID: nullableID( - head.NormalContextID, - ), CompactionContextID: nullableID(head.CompactionContextID), - ToolOutputID: nullableID( - head.PendingToolOutputID, - ), LimitOutputID: nullableID(head.PendingOutputLimitID), - ConfigID: nullableID(head.PendingConfigInputID), CheckpointID: nullableID(head.PendingCheckpointID)} -} - -func loadExecutionFacts( - ctx context.Context, +func loadExecutionFacts(ctx context.Context, db executiondb.DBTX, route AgentRoute, - head ExecutionHead, -) (ExecutionSnapshot, error) { - batch := executiondb.New(). - LoadExecutionFacts(ctx, db, []executiondb.LoadExecutionFactsParams{executionFactsParams(route, head)}) - var snapshot ExecutionSnapshot - var err error - batch.Query(func(_ int, rows []executiondb.LoadExecutionFactsRow, queryErr error) { - if queryErr != nil { - err = queryErr - return - } - snapshot, err = executionSnapshot(route, head, rows) - }) - return snapshot, errors.Join(err, batch.Close()) + head ExecutionHead) (ExecutionSnapshot, + error) { + batch := factsBatch{DBTX: db} + facts := batch.queue(ctx, route, head) + if err := batch.read(ctx); err != nil { + return ExecutionSnapshot{}, err + } + return executionSnapshot(route, head, *facts) } -func executionSnapshot( - route AgentRoute, - head ExecutionHead, - rows []executiondb.LoadExecutionFactsRow, -) (ExecutionSnapshot, error) { - view := ExecutionView{AgentID: route.AgentID, StopSequence: head.StopSequence, - MaxNormalInputSequence: head.MaxNormalInputSequence, MaxContextInputSequence: head.MaxContextInputSequence} - var now time.Time - for _, row := range rows { - if row.Kind != "scope" { - continue - } - if row.ID != route.AgentID || row.RootAgentID != route.RootAgentID || - row.TurnID != head.CurrentTurnID { - return ExecutionSnapshot{}, fmt.Errorf("%w: execution scope", ErrInvalidState) - } - view.State = AgentState(row.State) - now = row.DatabaseNow - view.Inputs = InputAvailability{Steering: row.Steering, Queued: row.Queued} - if row.TurnID != uuid.Nil { - view.Turn = &Turn{ID: row.TurnID, FirstOpeningSequence: row.FirstOpeningSequence, - LastOpeningSequence: row.LastOpeningSequence, FirstContentSequence: row.FirstContentSequence, - LatestSemantic: EventBoundary{Sequence: row.Sequence, Time: row.EventTime}} - } +func executionSnapshot(route AgentRoute, head ExecutionHead, facts executionFacts) (ExecutionSnapshot, error) { + scope := facts.scope + if scope.AgentID != route.AgentID || + scope.RootAgentID != route.RootAgentID || + valueOrZero(scope.TurnID) != head.CurrentTurnID { + return ExecutionSnapshot{}, fmt.Errorf("%w: execution scope", ErrInvalidState) } - if now.IsZero() { - return ExecutionSnapshot{}, fmt.Errorf("%w: missing agent", ErrInvalidState) + now := scope.DatabaseNow + view := ExecutionView{AgentID: route.AgentID, State: AgentState(scope.State), StopSequence: head.StopSequence, + MaxNormalInputSequence: head.MaxNormalInputSequence, MaxContextInputSequence: head.MaxContextInputSequence, + Inputs: InputAvailability{Steering: scope.Steering, Queued: scope.Queued}} + if scope.TurnID != nil { + view.Turn = &Turn{ID: *scope.TurnID, + FirstOpeningSequence: scope.FirstOpeningSequence, + LastOpeningSequence: scope.LastOpeningSequence, + FirstContentSequence: scope.FirstContentSequence, + LatestSemantic: EventBoundary{Sequence: scope.SemanticSequence, + Time: scope.SemanticTime}} + if len(scope.OpeningIds) > 0 { + view.Turn.InitialOpening = Opening{InputIDs: scope.OpeningIds, EventSequence: scope.OpeningSequence} + view.Turn.InitialReadyAt = scope.OpeningReadyAt + } } contexts := make(map[uuid.UUID]*ModelContext) - var batch *ToolBatch - for _, row := range rows { - switch row.Kind { - case "context": - contexts[row.ID] = &ModelContext{ID: row.ID, AgentID: route.AgentID, TurnID: row.TurnID, - Operation: Operation(row.Operation), Attempt: row.Attempt, InputEventSequence: row.Watermark, - SourceEventSequenceEnd: row.SourceEnd, State: ContextState(row.State), Recovery: RecoveryKind(row.Recovery), - RetryAt: row.RetryAt, Opening: Opening{InputIDs: row.OpeningIds, EventSequence: row.OpeningSequence}} - case "batch": - if !row.HasTools { - return ExecutionSnapshot{}, fmt.Errorf("%w: pending output has no tools", ErrInvalidState) - } - batch = &ToolBatch{HasIncomplete: row.Incomplete, HasRunnable: row.Runnable} - if !row.Incomplete { - if row.Calls == 0 || row.Calls != row.Results || row.Calls != row.Events { - return ExecutionSnapshot{}, fmt.Errorf("%w: incomplete result authority", ErrInvalidState) - } - batch.Completion = &BatchCompletion{LastResultSequence: row.Sequence, ReadyAt: row.EventTime} - } - case "config", "checkpoint": - work := &BoundaryWork{ID: row.ID, AgentID: route.AgentID, TurnID: row.TurnID, - Event: EventBoundary{Sequence: row.Sequence, Time: row.EventTime}, - Opening: Opening{InputIDs: row.OpeningIds, EventSequence: row.OpeningSequence}} - if row.Kind == "config" { - view.Config = work - } else { - view.Checkpoint = work - } - case "opening": - if view.Turn == nil || head.MaxContextInputSequence >= view.Turn.FirstOpeningSequence || - head.StopSequence > view.Turn.FirstOpeningSequence { - return ExecutionSnapshot{}, ErrInvalidState - } - if len(view.Turn.InitialOpening.InputIDs) == 0 { - view.Turn.InitialOpening.EventSequence = row.Sequence - view.Turn.InitialReadyAt = row.EventTime - } - view.Turn.InitialOpening.InputIDs = append(view.Turn.InitialOpening.InputIDs, row.ID) - if row.EventTime.Before(view.Turn.InitialReadyAt) { - view.Turn.InitialReadyAt = row.EventTime - } - } + for _, row := range facts.contexts { + contexts[row.ID] = &ModelContext{ID: row.ID, + AgentID: row.AgentID, + TurnID: row.TurnID, + Operation: Operation(row.OperationKind), + Attempt: row.AttemptNumber, + InputEventSequence: row.InputEventSequence, + SourceEventSequenceEnd: valueOrZero(row.SourceEventSequenceEnd), + State: ContextState(row.State), + Recovery: RecoveryKind(valueOrZero(row.RecoveryKind)), + RetryAt: row.RetryAt, + Opening: Opening{InputIDs: row.OpeningInputIds, EventSequence: row.OpeningEventSequence}} } view.NormalContext = contexts[head.NormalContextID] view.CompactionContext = contexts[head.CompactionContextID] - for _, row := range rows { - if row.Kind != "output" { + for _, row := range facts.configs { + view.Config = &BoundaryWork{ID: row.ID, AgentID: route.AgentID, TurnID: row.TurnID, + Event: EventBoundary{Sequence: row.Sequence, + Time: row.CreatedAt}, + Opening: Opening{InputIDs: row.OpeningInputIds, + EventSequence: row.OpeningEventSequence}} + } + for _, row := range facts.checkpoints { + view.Checkpoint = &BoundaryWork{ID: row.ID, AgentID: route.AgentID, TurnID: row.TurnID, + Event: EventBoundary{Sequence: row.Sequence, + Time: row.CreatedAt}, + Opening: Opening{InputIDs: row.OpeningInputIds, + EventSequence: row.OpeningEventSequence}} + } + for _, row := range facts.contexts { + if row.OutputID == nil { continue } - source := contexts[row.ContextID] + source := contexts[row.ID] if source == nil { return ExecutionSnapshot{}, fmt.Errorf("%w: missing output context", ErrInvalidState) } - output := ModelOutput{ - ID: row.ID, - EventID: row.EventID, - Event: EventBoundary{Sequence: row.Sequence, Time: row.EventTime}, + output := ModelOutput{ID: *row.OutputID, + EventID: valueOrZero(row.OutputEventID), Context: *source, - } - if row.ID == head.PendingToolOutputID && batch != nil { - batch.Output = output - view.ToolBatch = batch - } else if row.ID == head.PendingOutputLimitID && row.StopReason == "max_tokens" { + Event: EventBoundary{Sequence: valueOrZero(row.OutputSequence), + Time: valueOrZero(row.OutputCreatedAt)}} + if *row.OutputID == head.PendingToolOutputID { + b := facts.scope + if !b.BatchHasTools { + return ExecutionSnapshot{}, fmt.Errorf("%w: pending output has no tools", ErrInvalidState) + } + view.ToolBatch = &ToolBatch{Output: output, HasIncomplete: b.BatchIncomplete, HasRunnable: b.BatchRunnable} + if !b.BatchIncomplete { + if b.BatchCalls == 0 || b.BatchCalls != b.BatchResults || b.BatchCalls != b.BatchEvents { + return ExecutionSnapshot{}, fmt.Errorf("%w: incomplete result authority", ErrInvalidState) + } + view.ToolBatch.Completion = &BatchCompletion{LastResultSequence: b.BatchLastSequence, ReadyAt: b.BatchReadyAt} + } + } else if *row.OutputID == head.PendingOutputLimitID && valueOrZero(row.OutputStopReason) == "max_tokens" { view.OutputLimit = &output } else { return ExecutionSnapshot{}, fmt.Errorf("%w: pending output", ErrInvalidState) diff --git a/internal/storage/internal/agentexecution/load_count_integration_test.go b/internal/storage/internal/agentexecution/load_count_integration_test.go index c58d2707d..dc62683b0 100644 --- a/internal/storage/internal/agentexecution/load_count_integration_test.go +++ b/internal/storage/internal/agentexecution/load_count_integration_test.go @@ -15,7 +15,7 @@ type countingLoadTx struct { statements int batches int queries int - facts *pgx.QueuedQuery + facts []*pgx.QueuedQuery factLoads int } @@ -45,19 +45,21 @@ func (tx *countingLoadTx) QueryRow(ctx context.Context, sql string, args ...any) return tx.Tx.QueryRow(ctx, sql, args...) } -func CountOrdinaryLoad(ctx context.Context, h *Handle) (ExecutionSnapshot, int, error) { +func CountOrdinaryLoad(ctx context.Context, h *Handle) (ExecutionSnapshot, int, int, error) { tx := &countingLoadTx{Tx: h.unit.tx} h.unit.tx = tx defer func() { h.unit.tx = tx.Tx }() snapshot, err := h.LoadExecution(ctx) - return snapshot, tx.statements, err + return snapshot, tx.statements, tx.batches + tx.queries, err } func (tx *countingLoadTx) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { tx.statements += batch.Len() for _, query := range batch.QueuedQueries { + if strings.HasPrefix(query.SQL, "-- name: LoadExecution") { + tx.facts = append(tx.facts, query) + } if strings.HasPrefix(query.SQL, "-- name: LoadExecutionFacts ") { - tx.facts = query tx.factLoads++ } } @@ -97,9 +99,8 @@ func ExplainOrdinaryLoad(ctx context.Context, h *Handle) (string, error) { return "", err } var plan strings.Builder - for range 10 { - plan.Reset() - rows, err := tx.Tx.Query(ctx, "EXPLAIN (ANALYZE, BUFFERS, TIMING OFF) "+tx.facts.SQL, tx.facts.Arguments...) + for _, query := range tx.facts { + rows, err := tx.Tx.Query(ctx, "EXPLAIN (ANALYZE, BUFFERS, TIMING OFF) "+query.SQL, query.Arguments...) if err != nil { return "", err } diff --git a/internal/storage/internal/agentexecution/load_integration_test.go b/internal/storage/internal/agentexecution/load_integration_test.go index 62de7ad99..ad1e8f643 100644 --- a/internal/storage/internal/agentexecution/load_integration_test.go +++ b/internal/storage/internal/agentexecution/load_integration_test.go @@ -37,16 +37,19 @@ func TestOrdinaryExecutionLoadStatementBound(t *testing.T) { for name, id := range ids { t.Run(name, func(t *testing.T) { _, h := f.handle(t, id) - snapshot, count, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + snapshot, count, roundTrips, err := agentexecution.CountOrdinaryLoad(t.Context(), h) require.NoError(t, err) - require.Equal(t, 2, count) - _, count, err = agentexecution.CountOrdinaryLoad(t.Context(), h) + expectedStatements := map[string]int{"idle": 2, "initial": 2, "incomplete": 3, "complete": 3, "retry": 3} + require.Equal(t, expectedStatements[name], count) + require.Equal(t, 2, roundTrips) + _, count, roundTrips, err = agentexecution.CountOrdinaryLoad(t.Context(), h) require.NoError(t, err) expected := 0 if snapshot.Selection.Wait == agentexecution.WaitModelDeadline { expected = 1 } require.Equal(t, expected, count) + require.Equal(t, expected, roundTrips) plan, err := agentexecution.ExplainOrdinaryLoad(t.Context(), h) require.NoError(t, err) t.Log(plan) diff --git a/internal/storage/internal/agentexecution/locks.go b/internal/storage/internal/agentexecution/locks.go index 2ee876918..ce8e3616f 100644 --- a/internal/storage/internal/agentexecution/locks.go +++ b/internal/storage/internal/agentexecution/locks.go @@ -14,32 +14,19 @@ import ( "github.com/omnara-ai/omnara/internal/storage/storeerr" ) -type MutationAuthority interface{ mutationAuthority() } - type RuntimeAuthority struct { AgentID uuid.UUID RuntimeLockID uuid.UUID } -type IngressAuthority struct{} -type LifecycleAuthority struct{} -type RepairAuthority struct{} -type ExternalAuthority struct{} - -func (RuntimeAuthority) mutationAuthority() {} -func (IngressAuthority) mutationAuthority() {} -func (LifecycleAuthority) mutationAuthority() {} -func (RepairAuthority) mutationAuthority() {} -func (ExternalAuthority) mutationAuthority() {} - type MutationPlan struct { agents []AgentRoute - authority MutationAuthority + authority *RuntimeAuthority unit *Unit } func (u *Unit) PlanAgents( - ctx context.Context, refs []lifecyclelock.AgentRef, authority MutationAuthority, + ctx context.Context, refs []lifecyclelock.AgentRef, authority *RuntimeAuthority, ) (MutationPlan, error) { if err := u.active(); err != nil { return MutationPlan{}, err @@ -61,7 +48,7 @@ func (u *Unit) PlanAgents( return MutationPlan{}, err } plan.agents = append(plan.agents, AgentRoute{ - CellID: u.cell.id, ProjectID: ref.ProjectID, AgentID: ref.AgentID, RootAgentID: root, + ProjectID: ref.ProjectID, AgentID: ref.AgentID, RootAgentID: root, }) } return plan, nil @@ -80,19 +67,27 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo if err := u.active(); err != nil { return false, err } - switch plan.authority.(type) { - case RuntimeAuthority, IngressAuthority, LifecycleAuthority, RepairAuthority, ExternalAuthority: - default: - return false, errors.New("agent mutation authority is required") - } if plan.unit != u { return false, fmt.Errorf("%w: plan belongs to a different unit", ErrLockPlan) } + if len(u.handles) == 0 && len(u.expanded.agents) > 0 { + expanded, err := u.PlanAgents(ctx, u.expanded.agents, nil) + if err != nil { + return false, err + } + plan.agents = append(plan.agents, expanded.agents...) + } routes := slices.Clone(plan.agents) slices.SortFunc(routes, compareRoutes) routes = slices.Compact(routes) + var missing []lifecyclelock.AgentRef for _, route := range routes { - if route.CellID != u.cell.id || route.ProjectID == uuid.Nil || + if u.handles[route.AgentID] == nil { + missing = append(missing, lifecyclelock.AgentRef{ProjectID: route.ProjectID, AgentID: route.AgentID}) + } + } + for _, route := range routes { + if route.ProjectID == uuid.Nil || route.AgentID == uuid.Nil || route.RootAgentID == uuid.Nil { return false, fmt.Errorf("%w: invalid agent route", ErrLockPlan) } @@ -103,7 +98,7 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo continue } if u.runtimeLocked { - return false, fmt.Errorf("%w: agent acquired after runtime", ErrLockPlan) + return false, u.expandPlan(missing...) } for _, held := range u.handles { _, excluded := u.exclusiveProjects[route.ProjectID] @@ -111,7 +106,7 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo continue } if compareRoutes(route, held.route) < 0 { - return false, fmt.Errorf("%w: agent order", ErrLockPlan) + return false, u.expandPlan(missing...) } } } @@ -144,7 +139,7 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo } u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true, identity: &identity} } - if owned, ok := plan.authority.(RuntimeAuthority); ok { + if owned := plan.authority; owned != nil { h := u.handles[owned.AgentID] if h == nil || owned.RuntimeLockID == uuid.Nil { return false, storeerr.ErrRuntimeLockInactive @@ -165,7 +160,7 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo return true, nil } -func (u *Unit) LockAgentRefs(ctx context.Context, refs []lifecyclelock.AgentRef, authority MutationAuthority) error { +func (u *Unit) LockAgentRefs(ctx context.Context, refs []lifecyclelock.AgentRef, authority *RuntimeAuthority) error { plan, err := u.PlanAgents(ctx, refs, authority) if err != nil { return err @@ -209,13 +204,27 @@ func (u *Unit) ClaimAgent(ctx context.Context) (AgentRoute, error) { if len(u.handles) != 0 { return AgentRoute{}, ErrLockPlan } + if u.expanded.claimed != uuid.Nil { + if err := u.LockAgentRefs(ctx, u.expanded.agents, nil); err != nil { + return AgentRoute{}, err + } + if _, err := dbsqlc.New(u.DB()).LockClaimableAgentWakeup(ctx, + dbsqlc.LockClaimableAgentWakeupParams{AgentID: u.expanded.claimed}); err != nil { + return AgentRoute{}, err + } + u.claimedAgent = u.expanded.claimed + return u.handles[u.claimedAgent].route, nil + } wakeup, err := dbsqlc.New(u.DB()).ClaimNextAgentWakeup(ctx) if err != nil { return AgentRoute{}, err } - route := AgentRoute{CellID: u.cell.id, ProjectID: wakeup.ProjectID, + route := AgentRoute{ProjectID: wakeup.ProjectID, AgentID: wakeup.AgentID, RootAgentID: wakeup.RootAgentID} - u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} + if u.handles[route.AgentID] == nil { + u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} + } + u.claimedAgent = route.AgentID return route, nil } @@ -234,13 +243,13 @@ func (u *Unit) CreatedAgent(projectID, agentID, parentID uuid.UUID) error { } rootID = parent.route.RootAgentID } - route := AgentRoute{CellID: u.cell.id, ProjectID: projectID, AgentID: agentID, RootAgentID: rootID} + route := AgentRoute{ProjectID: projectID, AgentID: agentID, RootAgentID: rootID} u.handles[agentID] = &Handle{unit: u, route: route, valid: true, created: true} return nil } func (u *Unit) LockAgent( - ctx context.Context, arg dbsqlc.LockAgentInProjectParams, authority MutationAuthority, + ctx context.Context, arg dbsqlc.LockAgentInProjectParams, authority *RuntimeAuthority, ) (dbsqlc.LockAgentInProjectRow, error) { if err := u.LockAgentRefs( ctx, @@ -264,7 +273,7 @@ func (u *Unit) LockAgent( } func (u *Unit) PlanAgentFamily( - ctx context.Context, projectID, agentID uuid.UUID, authority MutationAuthority, + ctx context.Context, projectID, agentID uuid.UUID, authority *RuntimeAuthority, ) (MutationPlan, error) { if err := u.active(); err != nil { return MutationPlan{}, err @@ -278,9 +287,9 @@ func (u *Unit) PlanAgentFamily( if err != nil { return MutationPlan{}, err } - routes := []AgentRoute{{CellID: u.cell.id, ProjectID: projectID, AgentID: agentID, RootAgentID: row.RootAgentID}} + routes := []AgentRoute{{ProjectID: projectID, AgentID: agentID, RootAgentID: row.RootAgentID}} if row.ParentAgentID != nil { - routes = append(routes, AgentRoute{CellID: u.cell.id, ProjectID: projectID, + routes = append(routes, AgentRoute{ProjectID: projectID, AgentID: *row.ParentAgentID, RootAgentID: row.RootAgentID}) } return MutationPlan{unit: u, agents: routes, authority: authority}, nil diff --git a/internal/storage/internal/agentexecution/mutation.go b/internal/storage/internal/agentexecution/mutation.go index 83fcc3150..d0b08c3f7 100644 --- a/internal/storage/internal/agentexecution/mutation.go +++ b/internal/storage/internal/agentexecution/mutation.go @@ -60,7 +60,6 @@ func (m *executionMutation) publicationParams(route AgentRoute) executiondb.Publ InsertHead: m.insertHead, CurrentTurnID: nullableID(h.CurrentTurnID), StopSequence: h.StopSequence, - AnsweredThroughSequence: h.AnsweredThroughSequence, MaxNormalInputSequence: h.MaxNormalInputSequence, MaxContextInputSequence: h.MaxContextInputSequence, NormalContextID: nullableID(h.NormalContextID), @@ -77,7 +76,8 @@ func (m *executionMutation) publicationParams(route AgentRoute) executiondb.Publ func (u *Unit) publishExecution(ctx context.Context) error { var handles []*Handle - var loads []executiondb.LoadExecutionFactsParams + loads := factsBatch{DBTX: u.DB()} + var facts []*executionFacts var reloaded []*Handle for _, h := range u.orderedHandles() { if h.mutation == nil || !h.mutation.dirty { @@ -86,7 +86,7 @@ func (u *Unit) publishExecution(ctx context.Context) error { handles = append(handles, h) if h.mutation.loaded == nil && h.mutation.selectionChanged { reloaded = append(reloaded, h) - loads = append(loads, executionFactsParams(h.route, h.mutation.head)) + facts = append(facts, loads.queue(ctx, h.route, h.mutation.head)) } } if len(handles) == 0 { @@ -94,33 +94,19 @@ func (u *Unit) publishExecution(ctx context.Context) error { } var failure error q := executiondb.New() - if len(loads) > 0 { - batch := q.LoadExecutionFacts(ctx, u.DB(), loads) - batch.Query(func(i int, rows []executiondb.LoadExecutionFactsRow, err error) { - if failure != nil { - return - } - if err != nil { - failure = err - return - } - h := reloaded[i] - snapshot, err := executionSnapshot(h.route, h.mutation.head, rows) - if err != nil { - failure = err - return - } - h.mutation.head = snapshot.Head - }) - if err := errors.Join(failure, batch.Close()); err != nil { + if err := loads.read(ctx); err != nil { + return err + } + for i, h := range reloaded { + snapshot, err := executionSnapshot(h.route, h.mutation.head, *facts[i]) + if err != nil { return err } + h.mutation.head = snapshot.Head } writes := make([]executiondb.PublishExecutionParams, len(handles)) for i, h := range handles { writes[i] = h.mutation.publicationParams(h.route) - _, captured := u.webhooks[h.route.AgentID] - writes[i].CaptureWebhook = !captured } published := q.PublishExecution(ctx, u.DB(), writes) published.QueryRow(func(i int, row executiondb.PublishExecutionRow, err error) { @@ -135,14 +121,12 @@ func (u *Unit) publishExecution(ctx context.Context) error { failure = fmt.Errorf("%w: missing execution head", ErrInvalidState) return } - if writes[i].CaptureWebhook { - target := webhookTarget{orgID: row.WebhookOrgID} - if err := json.Unmarshal(row.WebhookEvents, &target.events); err != nil { - failure = err - return - } - u.webhooks[row.AgentID] = target + target := webhookTarget{orgID: row.WebhookOrgID} + if err := json.Unmarshal(row.WebhookEvents, &target.events); err != nil { + failure = err + return } + u.webhooks[row.AgentID] = target }) return errors.Join(failure, published.Close()) } diff --git a/internal/storage/internal/agentexecution/publication_integration_test.go b/internal/storage/internal/agentexecution/publication_integration_test.go index 8d703b219..d81dd659f 100644 --- a/internal/storage/internal/agentexecution/publication_integration_test.go +++ b/internal/storage/internal/agentexecution/publication_integration_test.go @@ -49,7 +49,11 @@ func TestExecutionCommitRoundTrips(t *testing.T) { if step == "model" { expected = 1 } - require.Equal(t, expected, statements) + expectedStatements := expected + if step == "tool" { + expectedStatements = 3 + } + require.Equal(t, expectedStatements, statements) require.Equal(t, expected, roundTrips) }) } @@ -63,19 +67,18 @@ func TestExecutionCommitRoundTrips(t *testing.T) { f.create(t, id, root) ids = append(ids, id) } - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(t.Context()) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(t.Context(), cell) require.NoError(t, err) t.Cleanup(func() { _ = u.Rollback(t.Context()) }) refs := make([]lifecyclelock.AgentRef, len(ids)) for i, id := range ids { refs[i] = lifecyclelock.AgentRef{ProjectID: f.project, AgentID: id} } - require.NoError(t, u.LockAgentRefs(t.Context(), refs, agentexecution.LifecycleAuthority{})) + require.NoError(t, u.LockAgentRefs(t.Context(), refs, nil)) for _, id := range ids { h, err := u.Agent( agentexecution.AgentRoute{ - CellID: "test", ProjectID: f.project, AgentID: id, RootAgentID: root, @@ -100,14 +103,13 @@ func TestExecutionCommitRoundTrips(t *testing.T) { require.Equal(t, size, heads) require.Equal(t, size, wakeups) - u, err = cell.Begin(t.Context()) + u, err = agentexecution.BeginTestUnit(t.Context(), cell) require.NoError(t, err) t.Cleanup(func() { _ = u.Rollback(t.Context()) }) - require.NoError(t, u.LockAgentRefs(t.Context(), refs, agentexecution.LifecycleAuthority{})) + require.NoError(t, u.LockAgentRefs(t.Context(), refs, nil)) routes := make([]agentexecution.AgentRoute, len(ids)) for i, id := range ids { routes[i] = agentexecution.AgentRoute{ - CellID: "test", ProjectID: f.project, AgentID: id, RootAgentID: root, diff --git a/internal/storage/internal/agentexecution/queries/load.sql b/internal/storage/internal/agentexecution/queries/load.sql index f50b5ccee..70c0183cf 100644 --- a/internal/storage/internal/agentexecution/queries/load.sql +++ b/internal/storage/internal/agentexecution/queries/load.sql @@ -1,7 +1,6 @@ -- name: LoadExecutionBase :one SELECT coalesce(h.agent_id,a.id)::uuid AS agent_id,h.current_turn_id, coalesce(h.stop_sequence,0)::bigint AS stop_sequence, - coalesce(h.answered_through_sequence,0)::bigint AS answered_through_sequence, coalesce(h.max_normal_input_sequence,0)::bigint AS max_normal_input_sequence, coalesce(h.max_context_input_sequence,0)::bigint AS max_context_input_sequence, h.normal_context_id,h.compaction_context_id,h.pending_tool_output_id,h.pending_output_limit_id, @@ -11,8 +10,24 @@ SELECT coalesce(h.agent_id,a.id)::uuid AS agent_id,h.current_turn_id, FROM agents a LEFT JOIN agent_execution_state h ON h.agent_id=a.id WHERE a.project_id=sqlc.arg(project_id) AND a.id=sqlc.arg(agent_id); --- name: LoadExecutionFacts :batchmany -WITH scope AS MATERIALIZED ( +-- name: LoadExecutionFacts :batchone +WITH batch AS ( +SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) AND c.model_output_id=sqlc.narg(tool_output_id)) AS has_tools, +EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) AND c.model_output_id=sqlc.narg(tool_output_id) AND c.state<>'completed') AS incomplete, +EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) AND c.model_output_id=sqlc.narg(tool_output_id) +AND (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable FROM (SELECT 1 WHERE sqlc.narg(tool_output_id)::uuid IS NOT NULL) present +), tool_batch AS ( +SELECT b.has_tools,b.incomplete,b.runnable,coalesce(r.calls,0)::bigint AS calls,coalesce(r.results,0)::bigint AS results, +coalesce(r.events,0)::bigint AS events,coalesce(r.last_sequence,0)::bigint AS last_sequence, +coalesce(r.ready_at,'epoch'::timestamptz)::timestamptz AS ready_at +FROM batch b LEFT JOIN LATERAL ( +SELECT count(c.id)::bigint AS calls,count(r.id)::bigint AS results,count(e.id)::bigint AS events, +max(e.sequence)::bigint AS last_sequence,max(r.completed_at)::timestamptz AS ready_at +FROM tool_calls c LEFT JOIN tool_call_results r ON r.agent_id=c.agent_id AND r.tool_call_id=c.id +LEFT JOIN agent_events e ON e.agent_id=r.agent_id AND e.tool_call_result_id=r.id +WHERE c.agent_id=sqlc.arg(agent_id) AND c.model_output_id=sqlc.narg(tool_output_id) AND NOT b.incomplete +) r ON NOT b.incomplete +), scope AS MATERIALIZED ( SELECT a.id,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now,t.id AS turn_id, coalesce(e.sequence,0)::bigint AS semantic_sequence,coalesce(e.created_at,'epoch'::timestamptz)::timestamptz AS semantic_time, EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' @@ -31,24 +46,11 @@ WITH scope AS MATERIALIZED ( WHERE x.agent_id=a.id AND x.turn_id=t.id AND x.is_opening_event ) o ON true WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) -), outputs AS MATERIALIZED ( - SELECT o.id,o.model_call_context_id,o.stop_reason,e.id AS event_id,e.sequence,e.created_at - FROM model_outputs o JOIN agent_events e ON e.agent_id=o.agent_id AND e.model_output_id=o.id - WHERE o.agent_id=sqlc.arg(agent_id) AND o.id IN (sqlc.narg(tool_output_id)::uuid,sqlc.narg(limit_output_id)::uuid) -), batch AS MATERIALIZED ( - SELECT EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) - AND c.model_output_id=sqlc.narg(tool_output_id)) AS has_tools, - EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) - AND c.model_output_id=sqlc.narg(tool_output_id) AND c.state<>'completed') AS incomplete, - EXISTS(SELECT 1 FROM tool_calls c WHERE c.agent_id=sqlc.arg(agent_id) - AND c.model_output_id=sqlc.narg(tool_output_id) AND - (c.state='awaiting_authorization' OR (c.state='ready' AND c.type IN ('built_in','mcp')))) AS runnable - WHERE sqlc.narg(tool_output_id)::uuid IS NOT NULL -), boundary AS MATERIALIZED ( +) , boundary AS MATERIALIZED ( SELECT greatest(coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=sqlc.arg(agent_id) AND e.event_kind='model_output' AND e.sequence<=(SELECT last_opening_sequence FROM scope) ORDER BY e.sequence DESC LIMIT 1),0),sqlc.arg(stop_sequence)::bigint)::bigint AS sequence - FROM scope WHERE scope.state='active' AND sqlc.arg(max_context_input_sequence)::bigint'event_webhook'->>'url','')<>'' - AND sqlc.arg(capture_webhook)::boolean ) w ON true; diff --git a/internal/storage/internal/agentexecution/queries/state.sql b/internal/storage/internal/agentexecution/queries/state.sql index 3d00becc3..447598868 100644 --- a/internal/storage/internal/agentexecution/queries/state.sql +++ b/internal/storage/internal/agentexecution/queries/state.sql @@ -7,8 +7,6 @@ SELECT coalesce((SELECT max(e.sequence) FROM agent_events e JOIN agent_inputs i ON i.agent_id=e.agent_id AND i.id=e.agent_input_id WHERE i.project_id=$1 AND i.agent_id=$2 AND i.input_kind='control' AND i.state='resolved' AND i.control_type='cancel_current'),0)::bigint AS stop_sequence, - coalesce((SELECT e.sequence FROM agent_events e WHERE e.agent_id=$2 AND e.event_kind='model_output' - ORDER BY e.sequence DESC LIMIT 1),0)::bigint AS output_sequence, coalesce((SELECT t.id FROM agent_turns t WHERE t.agent_id=$2 ORDER BY t.turn_sequence DESC LIMIT 1), '00000000-0000-0000-0000-000000000000'::uuid)::uuid AS turn_id; diff --git a/internal/storage/internal/agentexecution/reconstruction_integration_test.go b/internal/storage/internal/agentexecution/reconstruction_integration_test.go index fbdefc980..485b2d625 100644 --- a/internal/storage/internal/agentexecution/reconstruction_integration_test.go +++ b/internal/storage/internal/agentexecution/reconstruction_integration_test.go @@ -387,14 +387,14 @@ func (f executionFixture) handle( ) (*agentexecution.Unit, *agentexecution.Handle) { t.Helper() ctx := t.Context() - cell := agentexecution.NewCell("test", f.pool, nil, nil) - u, err := cell.Begin(ctx) + cell := agentexecution.NewCell(f.pool, nil, nil) + u, err := agentexecution.BeginTestUnit(ctx, cell) require.NoError(t, err) t.Cleanup(func() { _ = u.Rollback(context.Background()) }) plan, err := u.PlanAgents( ctx, []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: agent}}, - agentexecution.RepairAuthority{}, + nil, ) require.NoError(t, err) require.NoError(t, u.LockAgents(ctx, plan)) @@ -404,7 +404,7 @@ func (f executionFixture) handle( u.DB().QueryRow(ctx, `SELECT root_agent_id FROM agents WHERE id=$1`, agent).Scan(&root), ) h, err := u.Agent( - agentexecution.AgentRoute{CellID: cell.ID(), ProjectID: f.project, AgentID: agent, RootAgentID: root}, + agentexecution.AgentRoute{ProjectID: f.project, AgentID: agent, RootAgentID: root}, ) require.NoError(t, err) return u, h diff --git a/internal/storage/internal/agentexecution/repair.go b/internal/storage/internal/agentexecution/repair.go index bd6048e85..2d152e8ff 100644 --- a/internal/storage/internal/agentexecution/repair.go +++ b/internal/storage/internal/agentexecution/repair.go @@ -32,10 +32,9 @@ func (h *Handle) ReconstructExecution(ctx context.Context) (ExecutionSnapshot, e return ExecutionSnapshot{}, err } head := ExecutionHead{ - AgentID: l.route.AgentID, - CurrentTurnID: boundaries.TurnID, - StopSequence: boundaries.StopSequence, - AnsweredThroughSequence: max(boundaries.StopSequence, boundaries.OutputSequence), + AgentID: l.route.AgentID, + CurrentTurnID: boundaries.TurnID, + StopSequence: boundaries.StopSequence, } contexts, err := l.q.ReconstructExecutionContexts( ctx, diff --git a/internal/storage/internal/agentexecution/retry.go b/internal/storage/internal/agentexecution/retry.go new file mode 100644 index 000000000..cfa26f86b --- /dev/null +++ b/internal/storage/internal/agentexecution/retry.go @@ -0,0 +1,84 @@ +package agentexecution + +import ( + "context" + "errors" + "fmt" + "slices" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" + "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" +) + +type LockPlanError struct { + Agents []lifecyclelock.AgentRef + claimed uuid.UUID +} + +func (e *LockPlanError) Error() string { + return fmt.Sprintf("%s: required agents %v", ErrLockPlan, e.Agents) +} +func (e *LockPlanError) Unwrap() error { return ErrLockPlan } + +type lockPlanKey struct{} +type retryPlan struct { + agents []lifecyclelock.AgentRef + claimed uuid.UUID +} + +func Retry[T any](ctx context.Context, operation string, run func(context.Context) (T, error)) (T, error) { + if _, nested := ctx.Value(lockPlanKey{}).(retryPlan); nested { + return run(ctx) + } + return storeutil.RetryTransaction(ctx, operation, func() (T, error) { + plan := retryPlan{} + for attempt := 0; ; attempt++ { + result, err := run(context.WithValue(ctx, lockPlanKey{}, plan)) + var missing *LockPlanError + if !errors.As(err, &missing) || attempt == 7 || ctx.Err() != nil { + return result, err + } + previous := len(plan.agents) + plan.agents = mergeAgentRefs(plan.agents, missing.Agents) + if len(plan.agents) == previous { + return result, err + } + if missing.claimed != uuid.Nil { + plan.claimed = missing.claimed + } + } + }) +} + +func (u *Unit) expandPlan(refs ...lifecyclelock.AgentRef) error { + var held []lifecyclelock.AgentRef + for _, h := range u.handles { + if !h.created { + held = append(held, lifecyclelock.AgentRef{ProjectID: h.route.ProjectID, AgentID: h.route.AgentID}) + } + } + agents := mergeAgentRefs(held, refs) + u.aborted = true + return &LockPlanError{Agents: agents, claimed: u.claimedAgent} +} + +func mergeAgentRefs(a, b []lifecyclelock.AgentRef) []lifecyclelock.AgentRef { + seen := make(map[lifecyclelock.AgentRef]struct{}, len(a)+len(b)) + for _, refs := range [][]lifecyclelock.AgentRef{a, b} { + for _, ref := range refs { + seen[ref] = struct{}{} + } + } + result := make([]lifecyclelock.AgentRef, 0, len(seen)) + for ref := range seen { + result = append(result, ref) + } + slices.SortFunc(result, func(a, b lifecyclelock.AgentRef) int { + return compareRoutes(AgentRoute{ProjectID: a.ProjectID, + AgentID: a.AgentID}, + AgentRoute{ProjectID: b.ProjectID, + AgentID: b.AgentID}) + }) + return result +} diff --git a/internal/storage/internal/agentexecution/route.go b/internal/storage/internal/agentexecution/route.go index 234002b37..d552c6465 100644 --- a/internal/storage/internal/agentexecution/route.go +++ b/internal/storage/internal/agentexecution/route.go @@ -6,10 +6,7 @@ import ( "github.com/google/uuid" ) -type CellID string - type AgentRoute struct { - CellID CellID ProjectID uuid.UUID AgentID uuid.UUID RootAgentID uuid.UUID diff --git a/internal/storage/internal/agentexecution/schema_integration_test.go b/internal/storage/internal/agentexecution/schema_integration_test.go index 94227c35f..997708be0 100644 --- a/internal/storage/internal/agentexecution/schema_integration_test.go +++ b/internal/storage/internal/agentexecution/schema_integration_test.go @@ -327,3 +327,73 @@ func TestExecutionRootBackfillRejectsCycle(t *testing.T) { require.NoError(t, err) require.ErrorContains(t, applyExecutionCutover(t.Context(), f.pool), "invalid agent ancestry") } + +func TestExecutionOpeningMembersAreAgentScopedAndOrdered(t *testing.T) { + f := releasedExecutionFixture(t) + a, b := f.agent(t), f.agent(t) + _, first := f.opening(t, a, 1, uuid.Nil, "content") + turn, second := f.opening(t, a, 2, uuid.Nil, "content") + _, foreign := f.opening(t, b, 1, uuid.Nil, "content") + normal := f.modelContext(t, a, 2, "normal", nil, 1) + f.fail(t, normal, "compact") + compact := f.modelContext(t, a, 2, "compaction", new(int64(2)), 1) + f.checkpoint(t, a, turn, compact, 3, 2) + require.NoError(t, applyExecutionCutover(t.Context(), f.pool)) + for _, table := range []string{"model_call_contexts", "context_checkpoints", "agent_inputs"} { + for _, tc := range []struct { + name string + ids []uuid.UUID + valid bool + }{ + {"foreign", []uuid.UUID{foreign}, false}, + {"missing", []uuid.UUID{uuid.New()}, false}, + {"reversed", []uuid.UUID{second, first}, false}, + {"mixed", []uuid.UUID{first, foreign}, false}, + {"valid", []uuid.UUID{first, second}, true}, + } { + t.Run(table+"/"+tc.name, func(t *testing.T) { + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + switch table { + case "agent_inputs": + id, event := uuid.New(), uuid.New() + _, err = tx.Exec(t.Context(), `INSERT INTO agent_inputs( + id,project_id,agent_id,state,input_kind,delivery_mode,agent_config_id,queued_at) + VALUES($1,$2,$3,'received','config_change','immediate',$4,now())`, id, f.project, a, f.config) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), `INSERT INTO agent_events( + id,agent_id,turn_id,sequence,event_kind,agent_input_id,is_opening_event,idempotency_key,created_at) + VALUES($1,$2,$3,4,'agent_input',$4,true,'agent_input:'||($4::uuid)::text,now())`, event, a, turn, id) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), `UPDATE agent_inputs + SET state='resolved',admitted_event_id=$2,admitted_at=now(),resolved_at=now(), + opening_input_ids=$3,opening_event_sequence=1 WHERE id=$1`, id, event, tc.ids) + default: + source, id, ids := normal, uuid.New(), tc.ids + if table == "context_checkpoints" { + source, ids = compact, []uuid.UUID{first, second} + } + _, err = tx.Exec(t.Context(), `INSERT INTO model_call_contexts( + id,org_id,project_id,agent_id,operation_kind,attempt_number,agent_config_id,configured_model_revision_id, + input_event_sequence,source_event_sequence_end,runtime_lock_id,created_at,turn_id, + opening_input_ids,opening_event_sequence) + SELECT $1,org_id,project_id,agent_id,operation_kind,2,agent_config_id,configured_model_revision_id, + input_event_sequence,source_event_sequence_end,runtime_lock_id,now(),turn_id,$3,1 + FROM model_call_contexts WHERE id=$2`, id, source, ids) + if table == "context_checkpoints" { + require.NoError(t, err) + _, err = tx.Exec(t.Context(), `INSERT INTO context_checkpoints( + agent_id,summarized_through_event_sequence,producer_model_call_context_id,summary,created_at, + opening_input_ids,opening_event_sequence) VALUES($1,1,$2,'summary',now(),$3,1)`, a, id, tc.ids) + } + } + if tc.valid { + require.NoError(t, err) + } else { + require.ErrorContains(t, err, "opening inputs must belong to the agent in event sequence order") + } + }) + } + } +} diff --git a/internal/storage/internal/agentexecution/selection_load.go b/internal/storage/internal/agentexecution/selection_load.go index 89351d36f..877c0f02f 100644 --- a/internal/storage/internal/agentexecution/selection_load.go +++ b/internal/storage/internal/agentexecution/selection_load.go @@ -13,7 +13,6 @@ type ExecutionHead struct { AgentID uuid.UUID CurrentTurnID uuid.UUID StopSequence int64 - AnsweredThroughSequence int64 MaxNormalInputSequence int64 MaxContextInputSequence int64 NormalContextID uuid.UUID @@ -98,8 +97,10 @@ func (h *Handle) LoadExecution(ctx context.Context) (ExecutionSnapshot, error) { func executionHead(row executiondb.LoadExecutionBaseRow) ExecutionHead { return ExecutionHead{ - AgentID: row.AgentID, CurrentTurnID: valueOrZero(row.CurrentTurnID), StopSequence: row.StopSequence, - AnsweredThroughSequence: row.AnsweredThroughSequence, MaxNormalInputSequence: row.MaxNormalInputSequence, + AgentID: row.AgentID, + CurrentTurnID: valueOrZero(row.CurrentTurnID), + StopSequence: row.StopSequence, + MaxNormalInputSequence: row.MaxNormalInputSequence, MaxContextInputSequence: row.MaxContextInputSequence, NormalContextID: valueOrZero(row.NormalContextID), CompactionContextID: valueOrZero( row.CompactionContextID, diff --git a/internal/storage/internal/agentexecution/testing_integration.go b/internal/storage/internal/agentexecution/testing_integration.go new file mode 100644 index 000000000..6269c63f6 --- /dev/null +++ b/internal/storage/internal/agentexecution/testing_integration.go @@ -0,0 +1,9 @@ +//go:build integration + +package agentexecution + +import "context" + +func BeginTestUnit(ctx context.Context, cell *AgentCell) (*Unit, error) { + return Retry(ctx, "test unit", cell.Begin) +} diff --git a/internal/storage/internal/agentexecution/transition_count_integration_test.go b/internal/storage/internal/agentexecution/transition_count_integration_test.go index a6ddd2d2a..cc9523831 100644 --- a/internal/storage/internal/agentexecution/transition_count_integration_test.go +++ b/internal/storage/internal/agentexecution/transition_count_integration_test.go @@ -18,8 +18,8 @@ func TestExecutionTransitionStatementBounds(t *testing.T) { statements, facts, err := agentexecution.CountExecutionTransition(t.Context(), u, func() { run(h) }) require.NoError(t, err) t.Logf("%s statements=%d facts=%d", name, statements, facts) - limits := map[string]int{"receive": 6, "prepare_initial": 13, "accept_tools": 12, - "authorize": 6, "run_tool": 7, "complete_prepare": 16, "accept_end": 11} + limits := map[string]int{"receive": 7, "prepare_initial": 14, "accept_tools": 13, + "authorize": 6, "run_tool": 8, "complete_prepare": 17, "accept_end": 12} require.LessOrEqual(t, statements, limits[name], name) maximumFacts := 1 if name == "authorize" { diff --git a/internal/storage/internal/agentexecution/unit.go b/internal/storage/internal/agentexecution/unit.go index 48f8c524c..0ec9017b4 100644 --- a/internal/storage/internal/agentexecution/unit.go +++ b/internal/storage/internal/agentexecution/unit.go @@ -18,6 +18,8 @@ var ErrLockPlan = errors.New("agent execution lock plan must be restarted") type Unit struct { cell *AgentCell + expanded retryPlan + claimedAgent uuid.UUID tx pgx.Tx notifications *notifications.TxNotifications handles map[uuid.UUID]*Handle @@ -42,8 +44,6 @@ func (u *Unit) DB() unitDB { return unitDB{unit: u} } func (u *Unit) Notifications() *notifications.TxNotifications { return u.notifications } -func (u *Unit) CellID() CellID { return u.cell.id } - func (u *Unit) active() error { if u.closed { return ErrUnitClosed @@ -142,18 +142,18 @@ func (u *Unit) Savepoint(ctx context.Context, run func(*Unit) error) (err error) } type unitSnapshot struct { + claimedAgent uuid.UUID handles map[uuid.UUID]*Handle states map[uuid.UUID]Handle notifications *notifications.TxNotifications - webhooks map[uuid.UUID]webhookTarget exclusiveProjects map[uuid.UUID]struct{} runtimeLocked bool } func (u *Unit) snapshot() unitSnapshot { s := unitSnapshot{ - handles: maps.Clone(u.handles), states: make(map[uuid.UUID]Handle), - notifications: u.notifications.Clone(), webhooks: maps.Clone(u.webhooks), + claimedAgent: u.claimedAgent, handles: maps.Clone(u.handles), states: make(map[uuid.UUID]Handle), + notifications: u.notifications.Clone(), exclusiveProjects: maps.Clone(u.exclusiveProjects), runtimeLocked: u.runtimeLocked, } @@ -164,10 +164,6 @@ func (u *Unit) snapshot() unitSnapshot { } s.states[id] = state } - for id, target := range s.webhooks { - target.events = slices.Clone(target.events) - s.webhooks[id] = target - } return s } @@ -182,8 +178,8 @@ func (u *Unit) restore(snapshot unitSnapshot) { *h = snapshot.states[id] } u.notifications.Restore(snapshot.notifications) - u.webhooks = snapshot.webhooks u.exclusiveProjects = snapshot.exclusiveProjects u.runtimeLocked = snapshot.runtimeLocked + u.claimedAgent = snapshot.claimedAgent u.aborted = false } diff --git a/internal/storage/internal/agentexecution/unit_integration_test.go b/internal/storage/internal/agentexecution/unit_integration_test.go index ecb3db301..1cae41bd1 100644 --- a/internal/storage/internal/agentexecution/unit_integration_test.go +++ b/internal/storage/internal/agentexecution/unit_integration_test.go @@ -33,13 +33,13 @@ CREATE TABLE agents(id uuid PRIMARY KEY, project_id uuid NOT NULL, org_id uuid N CREATE TABLE unit_probe(id bigint GENERATED ALWAYS AS IDENTITY, agent_id uuid, phase text, labels text[]); CREATE TABLE deferred_probe(id integer UNIQUE DEFERRABLE INITIALLY DEFERRED);`) require.NoError(t, err) - return NewCell("test", pool, publisher, nil), pool + return NewCell(pool, publisher, nil), pool } func unitTestAgent(t *testing.T, u *Unit, project, agent uuid.UUID) *Handle { t.Helper() require.NoError(t, u.CreatedAgent(project, agent, uuid.Nil)) - h, err := u.Agent(AgentRoute{CellID: u.CellID(), ProjectID: project, AgentID: agent, RootAgentID: agent}) + h, err := u.Agent(AgentRoute{ProjectID: project, AgentID: agent, RootAgentID: agent}) require.NoError(t, err) return h } @@ -62,7 +62,7 @@ func TestUnitCommitLifecycle(t *testing.T) { published = append(published, intent) }}, ) - u, err := cell.Begin(ctx) + u, err := BeginTestUnit(ctx, cell) require.NoError(t, err) defer func() { _ = u.Rollback(ctx) }() _, raw := any(u.DB()).(pgx.Tx) @@ -104,7 +104,7 @@ func TestUnitSavepointSnapshots(t *testing.T) { published = append(published, intent) }}, ) - u, err := cell.Begin(ctx) + u, err := BeginTestUnit(ctx, cell) require.NoError(t, err) defer func() { _ = u.Rollback(ctx) }() project, agent, machine := uuid.New(), uuid.New(), uuid.New() @@ -115,15 +115,12 @@ func TestUnitSavepointSnapshots(t *testing.T) { pending.AddDaemonWork(machine) process := uuid.New() pending.AddDaemonProcessTermination(machine, process) - u.webhooks[agent] = webhookTarget{orgID: uuid.New(), events: []string{"model_output"}} - originalTarget := u.webhooks[agent] rejected := errors.New("reject inner") var escaped *Handle require.ErrorIs(t, u.Savepoint(ctx, func(u *Unit) error { mutation.head.StopSequence = 2 pending.AddDaemonWork(uuid.New()) pending.AddDaemonProcessTermination(machine, uuid.New()) - u.webhooks[agent] = webhookTarget{orgID: uuid.New()} delete(u.handles, agent) escaped = unitTestAgent(t, u, project, uuid.New()) escaped.mutation = &executionMutation{head: ExecutionHead{StopSequence: 3}} @@ -142,7 +139,6 @@ func TestUnitSavepointSnapshots(t *testing.T) { require.Same(t, h, u.handles[agent]) restored := h.mutation require.EqualValues(t, 1, restored.head.StopSequence) - require.Equal(t, originalTarget, u.webhooks[agent]) _, err = escaped.Route() require.ErrorIs(t, err, ErrUnitClosed) var unlocked bool @@ -179,7 +175,7 @@ func TestUnitFailuresSuppressPublication(t *testing.T) { t.Error("notification from failed transaction") }}, ) - u, err := cell.Begin(ctx) + u, err := BeginTestUnit(ctx, cell) require.NoError(t, err) defer func() { _ = u.Rollback(ctx) }() unitTestAgent(t, u, uuid.New(), uuid.New()) @@ -240,20 +236,22 @@ func TestUnitLockPlansAndSavepointLocks(t *testing.T) { org, ) require.NoError(t, err) - u, err := cell.Begin(ctx) + u, err := BeginTestUnit(ctx, cell) require.NoError(t, err) defer func() { _ = u.Rollback(ctx) }() refs := []lifecyclelock.AgentRef{{ProjectID: project, AgentID: higher}} - require.NoError(t, u.LockAgentRefs(ctx, refs, LifecycleAuthority{})) + require.NoError(t, u.LockAgentRefs(ctx, refs, nil)) require.Equal(t, lower, u.handles[higher].route.RootAgentID) - require.ErrorIs(t, u.LockAgentRefs( - ctx, - []lifecyclelock.AgentRef{{ProjectID: project, AgentID: lower}}, IngressAuthority{}), - ErrLockPlan, - ) - plan, err := u.PlanAgents(ctx, refs, LifecycleAuthority{}) + plan, err := u.PlanAgents(ctx, refs, nil) require.NoError(t, err) - other, err := NewCell("other", pool, nil, nil).Begin(ctx) + require.ErrorIs(t, + u.LockAgentRefs(ctx, + []lifecyclelock.AgentRef{{ProjectID: project, + AgentID: lower}}, + nil), + ErrLockPlan) + require.ErrorIs(t, u.active(), ErrUnitAborted) + other, err := BeginTestUnit(ctx, NewCell(pool, nil, nil)) require.NoError(t, err) defer func() { _ = other.Rollback(ctx) }() require.ErrorIs(t, other.LockAgents(ctx, plan), ErrLockPlan) @@ -261,7 +259,7 @@ func TestUnitLockPlansAndSavepointLocks(t *testing.T) { require.ErrorIs(t, other.Savepoint(ctx, func(u *Unit) error { return errors.Join(u.LockAgentRefs( ctx, - []lifecyclelock.AgentRef{{ProjectID: project, AgentID: lower}}, RepairAuthority{}), + []lifecyclelock.AgentRef{{ProjectID: project, AgentID: lower}}, nil), rollback, ) }), rollback) @@ -270,3 +268,195 @@ func TestUnitLockPlansAndSavepointLocks(t *testing.T) { _, err = check.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, lower) require.NoError(t, err) } + +func TestUnitExpandsLockPlanAndRollsBackAttempt(t *testing.T) { + ctx := t.Context() + cell, pool := unitTestCell(t, nil) + project, org := uuid.New(), uuid.New() + low, + high := uuid.MustParse("00000000-0000-4000-8000-000000000001"), + uuid.MustParse("ffffffff-ffff-4fff-8fff-ffffffffffff") + _, + err := pool.Exec(ctx, + `INSERT INTO agents(id,project_id,org_id,root_agent_id) VALUES($1,$3,$4,$1),($2,$3,$4,$2)`, + low, + high, + project, + org) + require.NoError(t, err) + for _, first := range []uuid.UUID{low, high} { + t.Run(first.String(), func(t *testing.T) { + second := low + if first == low { + second = high + } + attempts := 0 + _, err := Retry(ctx, "expand", func(ctx context.Context) (struct{}, error) { + attempts++ + u, err := BeginTestUnit(ctx, cell) + if err != nil { + return struct{}{}, err + } + defer func() { _ = u.Rollback(ctx) }() + if err = u.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: project, AgentID: first}}, nil); err != nil { + return struct{}{}, err + } + if _, err = u.DB().Exec(ctx, `INSERT INTO unit_probe(agent_id,phase) VALUES($1,'attempt')`, first); err != nil { + return struct{}{}, err + } + if _, err = u.Handle(project, second); err != nil { + return struct{}{}, err + } + return struct{}{}, u.Commit(ctx, "expanded") + }) + require.NoError(t, err) + require.Equal(t, 2, attempts) + var count int + require.NoError(t, pool.QueryRow(ctx, `SELECT count(*) FROM unit_probe WHERE agent_id=$1`, first).Scan(&count)) + require.Equal(t, 1, count) + }) + } +} + +func TestUnitLockPlanExpansionIsBounded(t *testing.T) { + attempts := 0 + _, err := Retry(t.Context(), "bounded", func(context.Context) (struct{}, error) { + attempts++ + return struct{}{}, &LockPlanError{Agents: []lifecyclelock.AgentRef{{ProjectID: uuid.New(), AgentID: uuid.New()}}} + }) + require.ErrorIs(t, err, ErrLockPlan) + require.Equal(t, 8, attempts) + attempts = 0 + ref := lifecyclelock.AgentRef{ProjectID: uuid.New(), AgentID: uuid.New()} + _, err = Retry(t.Context(), "no progress", func(context.Context) (struct{}, error) { + attempts++ + return struct{}{}, &LockPlanError{Agents: []lifecyclelock.AgentRef{ref}} + }) + require.ErrorIs(t, err, ErrLockPlan) + require.Equal(t, 2, attempts) +} + +func TestUnitExpandedClaimKeepsOriginalAgent(t *testing.T) { + ctx := t.Context() + cell, pool := unitTestCell(t, nil) + project, org, first, other, parent := uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New() + _, err := pool.Exec(ctx, ` +CREATE TABLE agent_wakeups(agent_id uuid PRIMARY KEY, ready_at timestamptz NOT NULL); +CREATE TABLE agent_runtime_locks(agent_id uuid PRIMARY KEY);`) + require.NoError(t, err) + _, err = pool.Exec(ctx, `INSERT INTO agents(id,project_id,org_id,root_agent_id) +SELECT id,$1,$2,id FROM unnest($3::uuid[]) id`, project, org, []uuid.UUID{first, other, parent}) + require.NoError(t, err) + _, err = pool.Exec(ctx, `INSERT INTO agent_wakeups VALUES($1,now()-interval '1 day'),($2,now())`, first, other) + require.NoError(t, err) + attempts := 0 + route, err := Retry(ctx, "expand claim", func(ctx context.Context) (AgentRoute, error) { + attempts++ + u, err := BeginTestUnit(ctx, cell) + if err != nil { + return AgentRoute{}, err + } + defer func() { _ = u.Rollback(ctx) }() + route, err := u.ClaimAgent(ctx) + if err != nil { + return route, err + } + if _, missing := u.Handle(project, parent); missing != nil { + require.NoError(t, u.Rollback(ctx)) + _, err = pool.Exec(ctx, `UPDATE agent_wakeups SET ready_at='epoch' WHERE agent_id=$1`, other) + require.NoError(t, err) + return AgentRoute{}, missing + } + return route, u.Commit(ctx, "expanded claim") + }) + require.NoError(t, err) + require.Equal(t, 2, attempts) + require.Equal(t, first, route.AgentID) +} + +func TestUnitExpandsCompleteRequestedPlan(t *testing.T) { + ctx := t.Context() + cell, pool := unitTestCell(t, nil) + project, org := uuid.New(), uuid.New() + first := uuid.MustParse("ffffffff-ffff-ffff-ffff-ffffffffffff") + ids := []uuid.UUID{first} + var refs []lifecyclelock.AgentRef + for i := byte(1); i <= 12; i++ { + id := uuid.UUID{15: i} + ids = append(ids, id) + refs = append(refs, lifecyclelock.AgentRef{ProjectID: project, AgentID: id}) + } + _, err := pool.Exec(ctx, `INSERT INTO agents(id,project_id,org_id,root_agent_id) +SELECT id,$1,$2,id FROM unnest($3::uuid[]) id`, project, org, ids) + require.NoError(t, err) + attempts := 0 + _, err = Retry(ctx, "complete expansion", func(ctx context.Context) (struct{}, error) { + attempts++ + u, err := BeginTestUnit(ctx, cell) + if err != nil { + return struct{}{}, err + } + defer func() { _ = u.Rollback(ctx) }() + if err = u.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: project, AgentID: first}}, nil); err != nil { + return struct{}{}, err + } + if err = u.LockAgentRefs(ctx, refs, nil); err != nil { + return struct{}{}, err + } + return struct{}{}, u.Commit(ctx, "complete expansion") + }) + require.NoError(t, err) + require.Equal(t, 2, attempts) +} + +func TestCellBeginRequiresRetryBoundary(t *testing.T) { + cell, _ := unitTestCell(t, nil) + unit, err := cell.Begin(t.Context()) + if unit != nil { + _ = unit.Rollback(t.Context()) + } + require.ErrorContains(t, err, "retry boundary") +} + +func TestUnitSavepointPropagatesLockPlan(t *testing.T) { + ctx := t.Context() + cell, pool := unitTestCell(t, nil) + project, org, first, second := uuid.New(), uuid.New(), uuid.New(), uuid.New() + _, err := pool.Exec(ctx, `INSERT INTO agents(id,project_id,org_id,root_agent_id) +SELECT id,$1,$2,id FROM unnest($3::uuid[]) id`, project, org, []uuid.UUID{first, second}) + require.NoError(t, err) + attempts := 0 + _, err = Retry(ctx, "savepoint expansion", func(ctx context.Context) (struct{}, error) { + attempts++ + unit, err := BeginTestUnit(ctx, cell) + if err != nil { + return struct{}{}, err + } + defer func() { _ = unit.Rollback(ctx) }() + if err = unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: project, AgentID: first}}, nil); err != nil { + return struct{}{}, err + } + err = unit.Savepoint(ctx, func(u *Unit) error { + _, err := u.DB().Exec(ctx, `INSERT INTO unit_probe(agent_id,phase) VALUES($1,'advance')`, first) + if err != nil { + return err + } + _, err = u.Handle(project, second) + return err + }) + if attempts == 1 { + var missing *LockPlanError + require.ErrorAs(t, err, &missing) + require.Contains(t, missing.Agents, lifecyclelock.AgentRef{ProjectID: project, AgentID: second}) + } + if err != nil { + return struct{}{}, err + } + return struct{}{}, unit.Commit(ctx, "expanded continuation") + }) + require.NoError(t, err) + require.Equal(t, 2, attempts) + var count int + require.NoError(t, pool.QueryRow(ctx, `SELECT count(*) FROM unit_probe WHERE agent_id=$1`, first).Scan(&count)) + require.Equal(t, 1, count) +} diff --git a/internal/storage/internal/agentexecution/webhooks.go b/internal/storage/internal/agentexecution/webhooks.go index 4f6006420..71041caff 100644 --- a/internal/storage/internal/agentexecution/webhooks.go +++ b/internal/storage/internal/agentexecution/webhooks.go @@ -17,11 +17,8 @@ type webhookTarget struct { events []string } -func (u *Unit) CaptureWebhookEligibility(ctx context.Context, agents ...uuid.UUID) error { +func (u *Unit) captureWebhookEligibility(ctx context.Context) error { ids := make(map[uuid.UUID]struct{}) - for _, id := range agents { - ids[id] = struct{}{} - } for id := range u.notifications.AgentEvents() { ids[id] = struct{}{} } @@ -66,7 +63,7 @@ func (u *Unit) enqueueWebhooks(ctx context.Context) error { if pending == nil || (len(pending.AgentEvents()) == 0 && len(pending.ToolCallUpdates()) == 0) { return nil } - if err := u.CaptureWebhookEligibility(ctx); err != nil { + if err := u.captureWebhookEligibility(ctx); err != nil { return err } q := dbsqlc.New(u.DB()) diff --git a/internal/storage/internal/dbsqlc/agent_inputs.sql.go b/internal/storage/internal/dbsqlc/agent_inputs.sql.go index d4f23763f..050cc5fc5 100644 --- a/internal/storage/internal/dbsqlc/agent_inputs.sql.go +++ b/internal/storage/internal/dbsqlc/agent_inputs.sql.go @@ -162,3 +162,23 @@ func (q *Queries) ListQueuedBacklogInputs(ctx context.Context, arg ListQueuedBac } return items, nil } + +const lockClaimableAgentWakeup = `-- name: LockClaimableAgentWakeup :one +SELECT wake.agent_id +FROM agent_wakeups wake JOIN agents agent ON agent.id=wake.agent_id +WHERE wake.agent_id=$1 AND agent.state<>'archived' + AND wake.ready_at<=statement_timestamp() + AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks runtime WHERE runtime.agent_id=wake.agent_id) +FOR UPDATE OF wake +` + +type LockClaimableAgentWakeupParams struct { + AgentID uuid.UUID +} + +func (q *Queries) LockClaimableAgentWakeup(ctx context.Context, arg LockClaimableAgentWakeupParams) (uuid.UUID, error) { + row := q.db.QueryRow(ctx, lockClaimableAgentWakeup, arg.AgentID) + var agent_id uuid.UUID + err := row.Scan(&agent_id) + return agent_id, err +} diff --git a/internal/storage/internal/dbsqlc/processes.sql.go b/internal/storage/internal/dbsqlc/processes.sql.go index 8dbbafd50..26b512ff2 100644 --- a/internal/storage/internal/dbsqlc/processes.sql.go +++ b/internal/storage/internal/dbsqlc/processes.sql.go @@ -1151,72 +1151,6 @@ func (q *Queries) ListActiveProcesses(ctx context.Context, arg ListActiveProcess return items, nil } -const listAgentsForExecutionRevoked = `-- name: ListAgentsForExecutionRevoked :many -SELECT agent.id -FROM agents agent -WHERE agent.project_id = $1 - AND ( - ($2::uuid IS NOT NULL AND agent.id = $2::uuid) - OR ( - $2::uuid IS NULL - AND ($3::uuid IS NOT NULL OR $4::uuid IS NOT NULL) - AND EXISTS ( - SELECT 1 - FROM agent_machine_bindings binding - JOIN project_machine_grants pmgrant ON pmgrant.project_id = binding.project_id - AND pmgrant.machine_id = binding.machine_id - WHERE binding.project_id = agent.project_id - AND binding.agent_id = agent.id - AND binding.state = 'attached' - AND ( - $3::uuid IS NULL - OR pmgrant.id = $3::uuid - ) - AND ( - $4::uuid IS NULL - OR ( - pmgrant.source_kind = 'pool' - AND pmgrant.project_machine_pool_grant_id = $4::uuid - ) - ) - ) - ) - ) -ORDER BY agent.id -` - -type ListAgentsForExecutionRevokedParams struct { - ProjectID uuid.UUID - AgentID *uuid.UUID - ProjectMachineGrantID *uuid.UUID - ProjectMachinePoolGrantID *uuid.UUID -} - -func (q *Queries) ListAgentsForExecutionRevoked(ctx context.Context, arg ListAgentsForExecutionRevokedParams) ([]uuid.UUID, error) { - rows, err := q.db.Query(ctx, listAgentsForExecutionRevoked, - arg.ProjectID, - arg.AgentID, - arg.ProjectMachineGrantID, - arg.ProjectMachinePoolGrantID, - ) - if err != nil { - return nil, err - } - defer rows.Close() - items := []uuid.UUID{} - for rows.Next() { - var id uuid.UUID - if err := rows.Scan(&id); err != nil { - return nil, err - } - items = append(items, id) - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const listDaemonProcessOffers = `-- name: ListDaemonProcessOffers :many WITH runtime AS MATERIALIZED ( SELECT runtime.org_id, runtime.machine_id diff --git a/internal/storage/orglifecycle/export_integration.go b/internal/storage/orglifecycle/export_integration.go index e88ffd7e5..b0067cff6 100644 --- a/internal/storage/orglifecycle/export_integration.go +++ b/internal/storage/orglifecycle/export_integration.go @@ -7,6 +7,7 @@ import ( "github.com/google/uuid" "github.com/omnara-ai/omnara/internal/storage/executionstore" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" ) func (s *Service) DeleteProjectOnceForIntegration( @@ -14,7 +15,11 @@ func (s *Service) DeleteProjectOnceForIntegration( orgID, projectID uuid.UUID, actor *executionstore.ActorParams, ) ([]executionstore.MachineRecord, error) { - return s.deleteProjectOnce(ctx, orgID, projectID, actor) + return agentexecution.Retry(ctx, + "test project deletion", + func(ctx context.Context) ([]executionstore.MachineRecord, error) { + return s.deleteProjectOnce(ctx, orgID, projectID, actor) + }) } func (s *Service) DeleteOrganizationOnceForIntegration( @@ -22,5 +27,9 @@ func (s *Service) DeleteOrganizationOnceForIntegration( orgID uuid.UUID, actor *executionstore.ActorParams, ) ([]executionstore.MachineRecord, error) { - return s.deleteOrganizationOnce(ctx, orgID, actor) + return agentexecution.Retry(ctx, + "test organization deletion", + func(ctx context.Context) ([]executionstore.MachineRecord, error) { + return s.deleteOrganizationOnce(ctx, orgID, actor) + }) } diff --git a/internal/storage/orglifecycle/orgs_projects.go b/internal/storage/orglifecycle/orgs_projects.go index a9e911de7..e6a075278 100644 --- a/internal/storage/orglifecycle/orgs_projects.go +++ b/internal/storage/orglifecycle/orgs_projects.go @@ -14,7 +14,6 @@ import ( "github.com/omnara-ai/omnara/internal/storage/internal/dbsqlc" "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/internal/skillops" - "github.com/omnara-ai/omnara/internal/storage/internal/storeutil" "github.com/omnara-ai/omnara/internal/storage/storeerr" ) @@ -276,7 +275,7 @@ func prelockProjectMachineLifecycleTx( if err := lifecyclelock.Machines(ctx, tx, machineRefs); err != nil { return nil, err } - if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { + if err := unit.LockAgentRefs(ctx, agentRefs, nil); err != nil { return nil, err } return agentIDs, nil @@ -374,7 +373,31 @@ func prelockOrganizationMachineLifecycleTx( if err := lifecyclelock.Machines(ctx, tx, machineRefs); err != nil { return organizationMachineLifecyclePlan{}, err } - if err := unit.LockAgentRefs(ctx, agentRefs, agentexecution.LifecycleAuthority{}); err != nil { + for _, poolID := range poolIDs { + rows, + err := q.ListMachinePoolAgentRefsForLifecycle(ctx, + dbsqlc.ListMachinePoolAgentRefsForLifecycleParams{OrgID: orgID, + MachinePoolID: poolID}) + if err != nil { + return organizationMachineLifecyclePlan{}, err + } + for _, row := range rows { + agentRefs = append(agentRefs, lifecyclelock.AgentRef{ProjectID: row.ProjectID, AgentID: row.AgentID}) + } + } + for _, machineID := range machineIDs { + rows, + err := q.ListMachineLifecycleTerminalAgentRefs(ctx, + dbsqlc.ListMachineLifecycleTerminalAgentRefsParams{OrgID: orgID, + MachineID: machineID}) + if err != nil { + return organizationMachineLifecyclePlan{}, err + } + for _, row := range rows { + agentRefs = append(agentRefs, lifecyclelock.AgentRef{ProjectID: row.ProjectID, AgentID: row.AgentID}) + } + } + if err := unit.LockAgentRefs(ctx, agentRefs, nil); err != nil { return organizationMachineLifecyclePlan{}, err } byoMachineIDs, err := q.ListActiveBYOMachineIDsForOrganizationDeletion( @@ -421,7 +444,7 @@ func (s *Service) DeleteProject( if err != nil { return nil, err } - return storeutil.RetryTransaction(ctx, "delete_project", func() ([]executionstore.MachineRecord, error) { + return agentexecution.Retry(ctx, "delete_project", func(ctx context.Context) ([]executionstore.MachineRecord, error) { return s.deleteProjectOnce(ctx, orgID, projectID, actor) }) } @@ -474,9 +497,6 @@ func (s *Service) deleteProjectOnce( if err != nil { return nil, err } - if err := unit.CaptureWebhookEligibility(ctx); err != nil { - return nil, err - } rows, err := q.DeleteProject(ctx, dbsqlc.DeleteProjectParams{OrgID: orgID, ID: projectID}) if err != nil { return nil, fmt.Errorf("delete project: %w", err) @@ -506,9 +526,12 @@ func (s *Service) DeleteOrganization( if err != nil { return nil, err } - return storeutil.RetryTransaction(ctx, "delete_organization", func() ([]executionstore.MachineRecord, error) { - return s.deleteOrganizationOnce(ctx, orgID, actor) - }) + return agentexecution.Retry(ctx, + "delete_organization", + func(ctx context.Context) ([]executionstore.MachineRecord, + error) { + return s.deleteOrganizationOnce(ctx, orgID, actor) + }) } //nolint:lll // Keeping generated query parameters inline makes the transaction's cascade order auditable. @@ -546,13 +569,6 @@ func (s *Service) deleteOrganizationOnce( if err != nil { return nil, err } - var webhookAgents []uuid.UUID - for _, ids := range plan.agentIDsByProject { - webhookAgents = append(webhookAgents, ids...) - } - if err := unit.CaptureWebhookEligibility(ctx, webhookAgents...); err != nil { - return nil, err - } rows, err := q.DeleteOrganization(ctx, dbsqlc.DeleteOrganizationParams{ID: orgID}) if err != nil { return nil, fmt.Errorf("delete organization: %w", err) diff --git a/internal/storage/orglifecycle/service.go b/internal/storage/orglifecycle/service.go index 70c0aff71..99d6c8f20 100644 --- a/internal/storage/orglifecycle/service.go +++ b/internal/storage/orglifecycle/service.go @@ -36,7 +36,7 @@ type Service struct { func New(pool *pgxpool.Pool, config Config) *Service { db := storeutil.WrapPool(pool) return &Service{ - cell: agentexecution.NewCell("primary", pool, config.PostCommitPublisher, nil), + cell: agentexecution.NewCell(pool, config.PostCommitPublisher, nil), pool: db, q: dbsqlc.New(db), blobs: config.Blobs, diff --git a/internal/storage/queries/agent_inputs.sql b/internal/storage/queries/agent_inputs.sql index 5babb0d4e..50a2dbfc9 100644 --- a/internal/storage/queries/agent_inputs.sql +++ b/internal/storage/queries/agent_inputs.sql @@ -32,6 +32,14 @@ locked_wake AS MATERIALIZED ( SELECT agent_id, project_id, root_agent_id FROM locked_wake; +-- name: LockClaimableAgentWakeup :one +SELECT wake.agent_id +FROM agent_wakeups wake JOIN agents agent ON agent.id=wake.agent_id +WHERE wake.agent_id=sqlc.arg(agent_id) AND agent.state<>'archived' + AND wake.ready_at<=statement_timestamp() + AND NOT EXISTS(SELECT 1 FROM agent_runtime_locks runtime WHERE runtime.agent_id=wake.agent_id) +FOR UPDATE OF wake; + -- name: ListQueuedBacklogInputs :many SELECT input.id, input.project_id, input.agent_id, input.state, input.input_rank, input.actor_id, input.input_kind, input.integration_target_id, coalesce(input.idempotency_scope, '') AS idempotency_scope, coalesce(input.input_idempotency_key, '') AS input_idempotency_key, input.queued_at, input.admitted_event_id, input.admitted_at, input.canceled_at, input.delivery_mode, coalesce(input.control_type, '') AS control_type, input.target_interaction_id, input.resolved_at, coalesce(input.rejected_reason, '') AS rejected_reason, input.metadata FROM agent_inputs input diff --git a/internal/storage/queries/processes.sql b/internal/storage/queries/processes.sql index d04042bfb..da920f72a 100644 --- a/internal/storage/queries/processes.sql +++ b/internal/storage/queries/processes.sql @@ -446,39 +446,6 @@ WHERE org_id = sqlc.arg(org_id) AND state IN ('starting', 'running') RETURNING id, org_id, project_id, agent_id, tool_call_id, runtime_lock_id, agent_machine_binding_id, machine_id, execution_granted_at, cwd, env, secret_env, timeout_seconds, initial_wait_ms, default_output_cursor, state, state_reason_code, state_reason_message, source_started_at, source_ended_at, state_changed_at, exit_code, exit_signal, created_at, updated_at, last_activity_at, execution_spec; --- name: ListAgentsForExecutionRevoked :many -SELECT agent.id -FROM agents agent -WHERE agent.project_id = sqlc.arg(project_id) - AND ( - (sqlc.narg(agent_id)::uuid IS NOT NULL AND agent.id = sqlc.narg(agent_id)::uuid) - OR ( - sqlc.narg(agent_id)::uuid IS NULL - AND (sqlc.narg(project_machine_grant_id)::uuid IS NOT NULL OR sqlc.narg(project_machine_pool_grant_id)::uuid IS NOT NULL) - AND EXISTS ( - SELECT 1 - FROM agent_machine_bindings binding - JOIN project_machine_grants pmgrant ON pmgrant.project_id = binding.project_id - AND pmgrant.machine_id = binding.machine_id - WHERE binding.project_id = agent.project_id - AND binding.agent_id = agent.id - AND binding.state = 'attached' - AND ( - sqlc.narg(project_machine_grant_id)::uuid IS NULL - OR pmgrant.id = sqlc.narg(project_machine_grant_id)::uuid - ) - AND ( - sqlc.narg(project_machine_pool_grant_id)::uuid IS NULL - OR ( - pmgrant.source_kind = 'pool' - AND pmgrant.project_machine_pool_grant_id = sqlc.narg(project_machine_pool_grant_id)::uuid - ) - ) - ) - ) - ) -ORDER BY agent.id; - -- name: ListProcessesForExecutionRevoked :many SELECT process.id, process.org_id, process.project_id, process.agent_id, process.tool_call_id, process.runtime_lock_id, process.agent_machine_binding_id, process.machine_id, process.execution_granted_at, process.cwd, process.env, process.secret_env, process.timeout_seconds, process.initial_wait_ms, process.default_output_cursor, process.state, process.state_reason_code, process.state_reason_message, process.source_started_at, process.source_ended_at, process.state_changed_at, process.exit_code, process.exit_signal, process.created_at, process.updated_at, process.last_activity_at, process.execution_spec FROM processes process diff --git a/migrations/000052_agent_execution_state.sql b/migrations/000052_agent_execution_state.sql index 0c566cf39..933dc1180 100644 --- a/migrations/000052_agent_execution_state.sql +++ b/migrations/000052_agent_execution_state.sql @@ -19,7 +19,6 @@ CREATE TABLE agent_execution_state ( agent_id uuid PRIMARY KEY REFERENCES agents(id) ON DELETE CASCADE, current_turn_id uuid, stop_sequence bigint NOT NULL DEFAULT 0 CHECK (stop_sequence >= 0), - answered_through_sequence bigint NOT NULL DEFAULT 0 CHECK (answered_through_sequence >= stop_sequence), max_normal_input_sequence bigint NOT NULL DEFAULT 0 CHECK (max_normal_input_sequence >= 0), max_context_input_sequence bigint NOT NULL DEFAULT 0 CHECK (max_context_input_sequence >= max_normal_input_sequence), normal_context_id uuid, @@ -174,10 +173,39 @@ CREATE TRIGGER model_call_contexts_lineage_immutable BEFORE UPDATE OF turn_id,opening_input_ids,opening_event_sequence ON model_call_contexts FOR EACH ROW EXECUTE FUNCTION enforce_execution_lineage_identity(); +-- +goose StatementBegin +CREATE FUNCTION enforce_execution_opening_members() RETURNS trigger LANGUAGE plpgsql AS $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM ( + SELECT member.ordinality,i.id,e.sequence, + lag(e.sequence) OVER (ORDER BY member.ordinality) AS previous_sequence + FROM unnest(NEW.opening_input_ids) WITH ORDINALITY member(id,ordinality) + LEFT JOIN agent_inputs i ON i.agent_id=NEW.agent_id AND i.id=member.id + AND i.input_kind='content' AND i.state='resolved' + LEFT JOIN agent_events e ON e.agent_id=i.agent_id AND e.id=i.admitted_event_id AND e.is_opening_event + ) ordered + WHERE id IS NULL OR sequence IS NULL OR sequence<=previous_sequence + OR (ordinality=1 AND sequence<>NEW.opening_event_sequence) + ) THEN + RAISE EXCEPTION 'opening inputs must belong to the agent in event sequence order' USING ERRCODE='23514'; + END IF; + RETURN NEW; +END; +$$; +-- +goose StatementEnd +CREATE TRIGGER model_call_contexts_opening_members BEFORE INSERT ON model_call_contexts + FOR EACH ROW EXECUTE FUNCTION enforce_execution_opening_members(); +CREATE TRIGGER context_checkpoints_opening_members BEFORE INSERT ON context_checkpoints + FOR EACH ROW EXECUTE FUNCTION enforce_execution_opening_members(); +CREATE TRIGGER agent_inputs_opening_members BEFORE INSERT OR UPDATE OF state,opening_input_ids,opening_event_sequence ON agent_inputs + FOR EACH ROW WHEN (NEW.input_kind='config_change' AND NEW.state='resolved') + EXECUTE FUNCTION enforce_execution_opening_members(); + INSERT INTO agent_execution_state ( - agent_id,current_turn_id,stop_sequence,answered_through_sequence,max_normal_input_sequence,max_context_input_sequence, + agent_id,current_turn_id,stop_sequence,max_normal_input_sequence,max_context_input_sequence, normal_context_id,compaction_context_id,turn_continuable,incomplete_tools,logical_ready_at) -SELECT a.id, t.id, coalesce(s.sequence,0), greatest(coalesce(s.sequence,0),coalesce(answer.sequence,0)), +SELECT a.id, t.id, coalesce(s.sequence,0), coalesce(w.normal_sequence,0), coalesce(w.context_sequence,0), n.id,c.id, a.state='active' AND (agent_has_incomplete_tool_batch(a.project_id,a.id) OR EXISTS(SELECT 1 FROM agent_next_model_work(a.project_id,a.id)) @@ -187,7 +215,6 @@ SELECT a.id, t.id, coalesce(s.sequence,0), greatest(coalesce(s.sequence,0),coale FROM agents a LEFT JOIN LATERAL (SELECT id FROM agent_turns WHERE agent_id=a.id ORDER BY turn_sequence DESC LIMIT 1) t ON true LEFT JOIN LATERAL (SELECT max(sequence) AS sequence FROM agent_stop_events WHERE agent_id=a.id AND project_id=a.project_id) s ON true -LEFT JOIN LATERAL (SELECT sequence FROM agent_events WHERE agent_id=a.id AND event_kind='model_output' ORDER BY sequence DESC LIMIT 1) answer ON true LEFT JOIN LATERAL (SELECT max(input_event_sequence) FILTER (WHERE operation_kind='normal') AS normal_sequence, max(input_event_sequence) AS context_sequence FROM model_call_contexts WHERE agent_id=a.id) w ON true LEFT JOIN LATERAL (SELECT id,input_event_sequence FROM model_call_contexts WHERE agent_id=a.id AND turn_id=t.id From 31f4d915516534c3814915b4133a9cc6b56281ba Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Mon, 12 Oct 2026 05:56:53 -0700 Subject: [PATCH 05/19] perf: drop the fused-acceptance savepoint and walk the ready index on claim - Run fused acceptance and advancement in one transaction. On a non-plan failure, roll back and accept alone in a fresh unit, as the commit-failure path already did. This removes the subtransaction that turned released lease rows into MultiXacts, saves two statements per fused step, and deletes the unit savepoint and snapshot machinery. - Stop anti-joining every runtime lease when choosing the next claim. Leased agents have no wakeup row, so the claim walks the ready index and stops at the first agent it can lock, with a single-row lease recheck. Co-Authored-By: Claude Opus 5.5 (1M context) --- internal/notifications/notifications.go | 35 ----- .../notifications_snapshot_test.go | 34 ---- .../event_webhooks_integration_test.go | 10 +- .../executionstore/model_work_transition.go | 51 +++--- .../commands_backlog_integration_test.go | 107 +++++++------ .../commands_custody_integration_test.go | 29 ++-- .../commands_family_integration_test.go | 24 +-- .../commands_integration_test.go | 142 +++++++++-------- .../commands_permission_integration_test.go | 18 ++- .../commands_processes_integration_test.go | 26 ++-- .../commands_runtime_integration_test.go | 31 ++-- .../storage/internal/agentexecution/handle.go | 3 +- .../storage/internal/agentexecution/locks.go | 6 +- .../internal/agentexecution/mutation.go | 13 -- .../agentexecution/schema_integration_test.go | 23 ++- .../storage/internal/agentexecution/unit.go | 92 +---------- .../agentexecution/unit_integration_test.go | 145 +----------------- .../internal/dbsqlc/agent_inputs.sql.go | 9 -- internal/storage/queries/agent_inputs.sql | 9 -- 19 files changed, 215 insertions(+), 592 deletions(-) delete mode 100644 internal/notifications/notifications_snapshot_test.go diff --git a/internal/notifications/notifications.go b/internal/notifications/notifications.go index 1fc8a14d8..3034126ef 100644 --- a/internal/notifications/notifications.go +++ b/internal/notifications/notifications.go @@ -5,8 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "maps" - "slices" "github.com/google/uuid" ) @@ -384,36 +382,3 @@ func (n *TxNotifications) Flush(ctx context.Context, publisher PostCommitPublish publisher.PublishPostCommit(ctx, intent) } } - -func (n *TxNotifications) Clone() *TxNotifications { - cloned := *n - cloned.daemonWorkByMachine = maps.Clone(n.daemonWorkByMachine) - cloned.runtimeEndedByID = maps.Clone(n.runtimeEndedByID) - cloned.processTerminationByMachine = make(map[uuid.UUID]map[uuid.UUID]struct{}, len(n.processTerminationByMachine)) - for id, processes := range n.processTerminationByMachine { - cloned.processTerminationByMachine[id] = maps.Clone(processes) - } - cloned.agentEventByID = make(map[uuid.UUID][]AgentEventReference, len(n.agentEventByID)) - for id, events := range n.agentEventByID { - cloned.agentEventByID[id] = slices.Clone(events) - } - cloned.toolCallUpdates = slices.Clone(n.toolCallUpdates) - for i, update := range cloned.toolCallUpdates { - if update.InteractionUpdate != nil { - interaction := *update.InteractionUpdate - cloned.toolCallUpdates[i].InteractionUpdate = &interaction - } - } - cloned.workerControls = slices.Clone(n.workerControls) - for i, control := range cloned.workerControls { - if control.Control.Cancel != nil { - cancel := *control.Control.Cancel - cloned.workerControls[i].Control.Cancel = &cancel - } - } - return &cloned -} - -func (n *TxNotifications) Restore(snapshot *TxNotifications) { - *n = *snapshot.Clone() -} diff --git a/internal/notifications/notifications_snapshot_test.go b/internal/notifications/notifications_snapshot_test.go deleted file mode 100644 index 6cb0fa7ae..000000000 --- a/internal/notifications/notifications_snapshot_test.go +++ /dev/null @@ -1,34 +0,0 @@ -package notifications - -import ( - "testing" - - "github.com/google/uuid" - "github.com/stretchr/testify/require" -) - -func TestTxNotificationsSnapshotOwnsNestedValues(t *testing.T) { - n := NewTxNotifications() - machine, process, runtime, agent, worker := uuid.New(), uuid.New(), uuid.New(), uuid.New(), uuid.New() - n.AddDaemonWork(machine) - n.AddDaemonProcessTermination(machine, process) - n.AddDaemonRuntimeEnded(runtime, machine, DaemonRuntimeEndReconnect) - n.AddAgentEvent(agent, 1, "agent_input") - n.AddToolCallUpdate(agent, uuid.New(), "waiting", &AgentInteractionUpdate{ID: uuid.New(), State: "open"}) - n.AddWorkerControlCancel(worker, agent, runtime) - snapshot := n.Clone() - expected := snapshot.Clone() - delete(n.daemonWorkByMachine, machine) - delete(n.processTerminationByMachine[machine], process) - delete(n.runtimeEndedByID, runtime) - n.agentEventByID[agent][0].Sequence = 2 - n.toolCallUpdates[0].InteractionUpdate.State = "resolved" - n.workerControls[0].Control.Cancel.RuntimeLockID = uuid.New() - require.Equal(t, expected, snapshot) - n.Restore(snapshot) - require.Equal(t, expected, n) - snapshot.toolCallUpdates[0].InteractionUpdate.State = "canceled" - snapshot.agentEventByID[agent][0].Sequence = 3 - delete(snapshot.processTerminationByMachine[machine], process) - require.Equal(t, expected, n) -} diff --git a/internal/storage/executionstore/event_webhooks_integration_test.go b/internal/storage/executionstore/event_webhooks_integration_test.go index e75382c11..6782ca7d0 100644 --- a/internal/storage/executionstore/event_webhooks_integration_test.go +++ b/internal/storage/executionstore/event_webhooks_integration_test.go @@ -336,7 +336,7 @@ func TestEventWebhookFiltersBeforeEnqueue(t *testing.T) { unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) defer func() { _ = unit.Rollback(ctx) }() - unit.Notifications().Restore(pending) + *unit.Notifications() = *pending require.NoError(t, unit.Commit(ctx, "filter test")) var sequences []int64 var toolUpdates int @@ -375,7 +375,7 @@ func TestEventWebhookClaimsRecoverAndExpire(t *testing.T) { unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) defer func() { _ = unit.Rollback(ctx) }() - unit.Notifications().Restore(pending) + *unit.Notifications() = *pending require.NoError(t, unit.Commit(ctx, "webhook test")) var attempts int var claimed bool @@ -537,7 +537,7 @@ event_webhook: unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) defer func() { _ = unit.Rollback(ctx) }() - unit.Notifications().Restore(pending) + *unit.Notifications() = *pending require.NoError(t, unit.Commit(ctx, "org limits test")) workers := []*executionstore.Store{fixture.Store.Execution(), newIntegrationStore(pool).Execution()} @@ -621,7 +621,7 @@ func TestEventWebhookCleanupIsBoundedAndPreservesUnexpiredDeliveries(t *testing. unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) defer func() { _ = unit.Rollback(ctx) }() - unit.Notifications().Restore(pending) + *unit.Notifications() = *pending require.NoError(t, unit.Commit(ctx, "cleanup test")) for i, toolID := range toolIDs { _, err := fixture.Store.pool.Exec(ctx, ` @@ -665,7 +665,7 @@ func TestEventWebhookDeletedOwnersHaveNoDeliveryTarget(t *testing.T) { unit, err := fixture.Store.Execution().IntegrationBeginUnit(ctx) require.NoError(t, err) defer func() { _ = unit.Rollback(ctx) }() - unit.Notifications().Restore(pending) + *unit.Notifications() = *pending require.NoError(t, unit.Commit(ctx, "deleted owner")) _, err = fixture.Store.pool.Exec(ctx, owner.query, owner.id) require.NoError(t, err) diff --git a/internal/storage/executionstore/model_work_transition.go b/internal/storage/executionstore/model_work_transition.go index 8531d5e00..1a58d9598 100644 --- a/internal/storage/executionstore/model_work_transition.go +++ b/internal/storage/executionstore/model_work_transition.go @@ -47,39 +47,26 @@ func (s *Store) acceptModelOutput(ctx context.Context, } result := OwnedAgentWorkTransition{} if advance && accepted.Created { - err = unit.Savepoint(ctx, func(_ *agentexecution.Unit) error { - work, err := h.Advance(ctx, output.RuntimeLockID, input.AllowModelWork, input.PrepareModel) - if err != nil { - return err - } - if err = applyAdmissionDestination(ctx, unit, output.ProjectID, output.AgentID, work.Admission); err != nil { - return err - } - shaped, err := shapeOwnedWork(ctx, unit, output.ProjectID, output.AgentID, work) - if err != nil { - return err - } - result = OwnedAgentWorkTransition{Work: shaped, Continued: !work.Released} - return nil - }) - } - if errors.Is(err, agentexecution.ErrLockPlan) { - return OwnedAgentWorkTransition{}, err - } - if !advance || !accepted.Created || err != nil { - result = OwnedAgentWorkTransition{Fallback: err} - if err = h.Release(ctx, output.RuntimeLockID); err != nil { - return result, err + var work agentexecution.OwnedWork + work, err = h.Advance(ctx, output.RuntimeLockID, input.AllowModelWork, input.PrepareModel) + if err == nil { + err = applyAdmissionDestination(ctx, unit, output.ProjectID, output.AgentID, work.Admission) } - } - if err = unit.Commit(ctx, "accept model output and advance"); err != nil { - if !advance || errors.Is(err, agentexecution.ErrLockPlan) { - return result, err + if err == nil { + result.Work, err = shapeOwnedWork(ctx, unit, output.ProjectID, output.AgentID, work) + result.Continued = !work.Released } - _ = unit.Rollback(ctx) - fallback, acceptErr := s.acceptModelOutput(ctx, input, false) - fallback.Fallback = errors.Join(result.Fallback, err, fallback.Fallback) - return fallback, acceptErr + } else { + err = h.Release(ctx, output.RuntimeLockID) + } + if err == nil { + err = unit.Commit(ctx, "accept model output and advance") + } + if err == nil || !advance || errors.Is(err, agentexecution.ErrLockPlan) { + return result, err } - return result, nil + _ = unit.Rollback(ctx) + fallback, acceptErr := s.acceptModelOutput(ctx, input, false) + fallback.Fallback = errors.Join(err, fallback.Fallback) + return fallback, acceptErr } diff --git a/internal/storage/internal/agentexecution/commands_backlog_integration_test.go b/internal/storage/internal/agentexecution/commands_backlog_integration_test.go index 633a6d752..f333c689d 100644 --- a/internal/storage/internal/agentexecution/commands_backlog_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_backlog_integration_test.go @@ -3,7 +3,6 @@ package agentexecution_test import ( - "errors" "testing" "github.com/google/uuid" @@ -33,8 +32,6 @@ func readCommandBacklog(t *testing.T, u *agentexecution.Unit, a uuid.UUID) []bac func TestExecutionBacklogOrderingAndReplay(t *testing.T) { f, a, _ := commandFixture(t) - u, h := f.handle(t, a) - rollback := errors.New("restore backlog") for _, dense := range []bool{false, true} { for _, position := range []string{"front", "back", "before", "after"} { name := position @@ -42,47 +39,45 @@ func TestExecutionBacklogOrderingAndReplay(t *testing.T) { name += "/rebalance" } t.Run(name, func(t *testing.T) { - err := u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - ids := make([]uuid.UUID, 3) - for i := range ids { - ids[i] = receiveCommand(t, h, "queued", "input").ID - if dense { - _, err := u.DB().Exec(t.Context(), - "UPDATE agent_inputs SET input_rank=$2 WHERE agent_id=$1 AND id=$3", a, i+1, ids[i]) - require.NoError(t, err) - } - } - change := agentexecution.BacklogChange{ - ID: ids[2], - DeliveryMode: "queued", - Position: position, - } - if position == "before" || position == "after" { - change.BeforeID = ids[0] - } - _, err := h.ChangeBacklog(t.Context(), change) - require.NoError(t, err) - expected := readCommandBacklog(t, u, a) - got := make([]uuid.UUID, len(expected)) - for i, row := range expected { - got[i] = row.ID - } - order := []uuid.UUID{ids[0], ids[1], ids[2]} - if position == "front" || position == "before" { - order = []uuid.UUID{ids[2], ids[0], ids[1]} - } - if position == "after" { - order = []uuid.UUID{ids[0], ids[2], ids[1]} + u, h := f.handle(t, a) + ids := make([]uuid.UUID, 3) + for i := range ids { + ids[i] = receiveCommand(t, h, "queued", "input").ID + if dense { + _, err := u.DB().Exec(t.Context(), + "UPDATE agent_inputs SET input_rank=$2 WHERE agent_id=$1 AND id=$3", a, i+1, ids[i]) + require.NoError(t, err) } - require.Equal(t, order, got) - assertCommandState(t, u, h) - changed, err := h.ChangeBacklog(t.Context(), change) - require.NoError(t, err) - require.False(t, changed) - require.Equal(t, expected, readCommandBacklog(t, u, a)) - return rollback - }) - require.ErrorIs(t, err, rollback) + } + change := agentexecution.BacklogChange{ + ID: ids[2], + DeliveryMode: "queued", + Position: position, + } + if position == "before" || position == "after" { + change.BeforeID = ids[0] + } + _, err := h.ChangeBacklog(t.Context(), change) + require.NoError(t, err) + expected := readCommandBacklog(t, u, a) + got := make([]uuid.UUID, len(expected)) + for i, row := range expected { + got[i] = row.ID + } + order := []uuid.UUID{ids[0], ids[1], ids[2]} + if position == "front" || position == "before" { + order = []uuid.UUID{ids[2], ids[0], ids[1]} + } + if position == "after" { + order = []uuid.UUID{ids[0], ids[2], ids[1]} + } + require.Equal(t, order, got) + assertCommandState(t, u, h) + changed, err := h.ChangeBacklog(t.Context(), change) + require.NoError(t, err) + require.False(t, changed) + require.Equal(t, expected, readCommandBacklog(t, u, a)) + require.NoError(t, u.Rollback(t.Context())) }) } } @@ -118,13 +113,15 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { require.NoError(t, err) require.True(t, changed) assertCommandState(t, u, h) - err = u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.ChangeBacklog( - t.Context(), - agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "queued"}, - ) - return err - }) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.ChangeBacklog( + t.Context(), + agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "queued"}, + ) + + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) require.ErrorIs(t, err, storeerr.ErrStateTransitionConflict) _, err = h.AdmitInputs(t.Context()) require.NoError(t, err) @@ -136,10 +133,12 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { changed, err = h.ChangeBacklog(t.Context(), cancel) require.NoError(t, err) require.True(t, changed) - err = u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.ChangeBacklog(t.Context(), cancel) - return err - }) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.ChangeBacklog(t.Context(), cancel) + + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) require.ErrorIs(t, err, storeerr.ErrStateTransitionConflict) assertCommandState(t, u, h) } diff --git a/internal/storage/internal/agentexecution/commands_custody_integration_test.go b/internal/storage/internal/agentexecution/commands_custody_integration_test.go index a529b5e34..a1f92b7e5 100644 --- a/internal/storage/internal/agentexecution/commands_custody_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_custody_integration_test.go @@ -61,14 +61,12 @@ func TestExecutionToolCustody(t *testing.T) { require.NoError(t, err) require.True(t, changed) assertCommandState(t, u, h) - require.ErrorIs( - t, - u.Savepoint( - t.Context(), - func(_ *agentexecution.Unit) error { _, err := h.RunTool(t.Context(), ref); return err }, - ), - storeerr.ErrToolCallInProgress, - ) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.RunTool(t.Context(), ref) + require.ErrorIs(t, err, storeerr.ErrToolCallInProgress) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) changed, err = h.RequeueTool(t.Context(), ref) require.NoError(t, err) require.True(t, changed) @@ -98,15 +96,12 @@ func TestExecutionToolCustody(t *testing.T) { require.NoError(t, err) require.False(t, replay.Changed) require.Equal(t, result.ResultID, replay.ResultID) - require.ErrorIs( - t, - u.Savepoint( - t.Context(), - func(_ *agentexecution.Unit) error { content[0].Text = "other"; _, err := complete(); return err }, - ), - storeerr.ErrIdempotencyConflict, - ) - require.NoError(t, u.Commit(t.Context(), "tool custody")) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + content[0].Text = "other" + _, err = complete() + require.ErrorIs(t, err, storeerr.ErrIdempotencyConflict) + require.NoError(t, u.Rollback(t.Context())) }) } } diff --git a/internal/storage/internal/agentexecution/commands_family_integration_test.go b/internal/storage/internal/agentexecution/commands_family_integration_test.go index 048271ce5..e59758246 100644 --- a/internal/storage/internal/agentexecution/commands_family_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_family_integration_test.go @@ -4,7 +4,6 @@ package agentexecution_test import ( "context" - "errors" "testing" "time" @@ -18,7 +17,7 @@ import ( "github.com/stretchr/testify/require" ) -func TestExecutionCreatedAgentPublicationAndSavepoint(t *testing.T) { +func TestExecutionCreatedAgentPublication(t *testing.T) { f := migratedExecutionFixture(t) cell := agentexecution.NewCell(f.pool, nil, nil) u, err := agentexecution.BeginTestUnit(t.Context(), cell) @@ -29,8 +28,6 @@ func TestExecutionCreatedAgentPublicationAndSavepoint(t *testing.T) { agentexecution.CreateAgentInput{ID: uuid.New(), ProjectID: f.project, ConfigID: f.config}, ) require.NoError(t, err) - route, err := h.Route() - require.NoError(t, err) receiveCommand(t, h, "queued", "created") assertCommandState(t, u, h) var heads int @@ -39,25 +36,6 @@ func TestExecutionCreatedAgentPublicationAndSavepoint(t *testing.T) { u.DB().QueryRow(t.Context(), `SELECT count(*) FROM agent_execution_state`).Scan(&heads), ) require.Zero(t, heads) - sentinel := errors.New("rollback child") - var escaped *agentexecution.Handle - require.ErrorIs(t, u.Savepoint(t.Context(), func(u *agentexecution.Unit) error { - escaped, err = u.CreateAgent( - t.Context(), - agentexecution.CreateAgentInput{ - ID: uuid.New(), - ProjectID: f.project, - ParentID: route.AgentID, - ConfigID: f.config, - SubagentKey: "child", - }, - ) - require.NoError(t, err) - receiveCommand(t, escaped, "queued", "discarded") - return sentinel - }), sentinel) - _, err = escaped.Route() - require.ErrorIs(t, err, agentexecution.ErrUnitClosed) statements, roundTrips, err := agentexecution.CountCommit(t.Context(), u) require.NoError(t, err) require.Equal(t, 1, statements) diff --git a/internal/storage/internal/agentexecution/commands_integration_test.go b/internal/storage/internal/agentexecution/commands_integration_test.go index 819622f73..8d561fff0 100644 --- a/internal/storage/internal/agentexecution/commands_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_integration_test.go @@ -5,7 +5,7 @@ package agentexecution_test import ( "context" "encoding/json" - "errors" + "slices" "testing" "time" @@ -178,18 +178,18 @@ func TestExecutionSemanticCommands(t *testing.T) { RuntimeLockID: lease, Response: commandResponse(modelenvelope.StopReasonMaxTokens), } - beforeOutput := u.Notifications().Clone() + beforeOutput := slices.Clone(u.Notifications().AgentEvents()[a]) accepted, err := h.AcceptOutput(t.Context(), input) require.NoError(t, err) require.True(t, accepted.Created) assertCommandState(t, u, h) - afterOutput := u.Notifications().Clone() + afterOutput := slices.Clone(u.Notifications().AgentEvents()[a]) require.NotEqual(t, beforeOutput, afterOutput) replayed, err := h.AcceptOutput(t.Context(), input) require.NoError(t, err) require.False(t, replayed.Created) require.Equal(t, accepted.ID, replayed.ID) - require.Equal(t, afterOutput, u.Notifications().Clone()) + require.Equal(t, afterOutput, slices.Clone(u.Notifications().AgentEvents()[a])) continuation := prepareCommand(t, h, lease) assertCommandState(t, u, h) _, err = h.AcceptOutput( @@ -203,7 +203,7 @@ func TestExecutionSemanticCommands(t *testing.T) { require.NoError(t, u.Commit(t.Context(), "semantic completion")) } -func TestExecutionCommandSavepointAndReplay(t *testing.T) { +func TestExecutionCommandRollbackAndReplay(t *testing.T) { f, a, _ := commandFixture(t) u, h := f.handle(t, a) input := agentexecution.ReceiveContentInput{ @@ -216,17 +216,14 @@ func TestExecutionCommandSavepointAndReplay(t *testing.T) { first, err := h.ReceiveContent(t.Context(), input) require.NoError(t, err) expected := assertCommandState(t, u, h) - sentinel := errors.New("rollback") - require.ErrorIs(t, u.Savepoint(t.Context(), func(u *agentexecution.Unit) error { - receiveCommand(t, h, "steering", "later") - _, err := h.AdmitInputs(t.Context()) - require.NoError(t, err) - assertCommandState(t, u, h) - return sentinel - }), sentinel) - _, count, _, err := agentexecution.CountOrdinaryLoad(t.Context(), h) + require.NoError(t, u.Commit(t.Context(), "input")) + u, h = f.handle(t, a) + receiveCommand(t, h, "steering", "later") + _, err = h.AdmitInputs(t.Context()) require.NoError(t, err) - require.Zero(t, count) + assertCommandState(t, u, h) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) got := assertCommandState(t, u, h) require.Equal(t, expected.Head.CurrentTurnID, got.Head.CurrentTurnID) require.Equal(t, agentexecution.AdmitOneQueued, got.Selection.Admission) @@ -234,12 +231,12 @@ func TestExecutionCommandSavepointAndReplay(t *testing.T) { require.NoError(t, err) require.False(t, replay.Created) require.Equal(t, first.ID, replay.ID) - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - input.Content[0].Text = "different" - _, err := h.ReceiveContent(t.Context(), input) - return err - }), storeerr.ErrIdempotencyConflict) - require.NoError(t, u.Commit(t.Context(), "savepoint replay")) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + input.Content[0].Text = "different" + _, err = h.ReceiveContent(t.Context(), input) + require.ErrorIs(t, err, storeerr.ErrIdempotencyConflict) + require.NoError(t, u.Rollback(t.Context())) u, h = f.handle(t, a) input.Content[0].Text = "one" replay, err = h.ReceiveContent(t.Context(), input) @@ -477,16 +474,18 @@ func TestExecutionConfigAndCheckpointBoundaries(t *testing.T) { replay, err := h.ActivateConfig(t.Context(), cfg) require.NoError(t, err) require.False(t, replay.Created) - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.PrepareModel( - t.Context(), - agentexecution.PrepareModelInput{ - RuntimeLockID: lease, - Selected: *selected.Selection.Model, - }, - ) - return err - }), storeerr.ErrAgentNotAdvanceable) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.PrepareModel( + t.Context(), + agentexecution.PrepareModelInput{ + RuntimeLockID: lease, + Selected: *selected.Selection.Model, + }, + ) + require.ErrorIs(t, err, storeerr.ErrAgentNotAdvanceable) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) } normal := prepareCommand(t, h, lease) if boundary == "steering" { @@ -633,17 +632,18 @@ func TestExecutionCheckpointRejectsCutBatch(t *testing.T) { ContextID: next.Context.ID, RuntimeLockID: lease, ErrorKind: "context_length", ErrorMessage: "compact"}}, ) require.NoError(t, err) - err = u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.PublishCheckpoint( - t.Context(), - agentexecution.PublishCheckpointInput{ContextID: compact.Model.Context.ID, - RuntimeLockID: lease, - Summary: "unsafe", - Evidence: agentexecution.ModelEvidence{APIFormat: "openai-responses", - APIVariant: "default"}}, - ) - return err - }) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.PublishCheckpoint( + t.Context(), + agentexecution.PublishCheckpointInput{ContextID: compact.Model.Context.ID, + RuntimeLockID: lease, + Summary: "unsafe", + Evidence: agentexecution.ModelEvidence{APIFormat: "openai-responses", + APIVariant: "default"}}, + ) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) require.ErrorContains(t, err, "boundary is unsafe") assertCommandState(t, u, h) var count int @@ -862,26 +862,27 @@ func TestExecutionOutputReplayRejectsConflictingEvidence(t *testing.T) { accepted, err := h.AcceptOutput(t.Context(), input) require.NoError(t, err) require.True(t, accepted.Created) + require.NoError(t, u.Commit(t.Context(), "accepted output")) for _, variant := range []string{"usage", "metadata", "content"} { t.Run(variant, func(t *testing.T) { - err := u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - conflict := input - switch variant { - case "usage": - conflict.Response.Normalized.Usage.OutputTokens++ - case "metadata": - conflict.Response.ProviderMetadata.OpenRouter.Provider = "different" - case "content": - conflict.Response.Normalized.Content = []modelenvelope.ResponsePart{ - {Type: "text", Text: "different"}, - } + u, h := f.handle(t, a) + conflict := input + switch variant { + case "usage": + conflict.Response.Normalized.Usage.OutputTokens++ + case "metadata": + conflict.Response.ProviderMetadata.OpenRouter.Provider = "different" + case "content": + conflict.Response.Normalized.Content = []modelenvelope.ResponsePart{ + {Type: "text", Text: "different"}, } - _, err := h.AcceptOutput(t.Context(), conflict) - return err - }) + } + _, err := h.AcceptOutput(t.Context(), conflict) + require.NoError(t, u.Rollback(t.Context())) require.ErrorIs(t, err, storeerr.ErrIdempotencyConflict) }) } + u, h = f.handle(t, a) replay, err := h.AcceptOutput(t.Context(), input) require.NoError(t, err) require.False(t, replay.Created) @@ -992,25 +993,22 @@ func TestExecutionArchiveHistoryDoesNotDependOnReads(t *testing.T) { _, err := h.AdmitInputs(t.Context()) require.NoError(t, err) require.NoError(t, u.Commit(t.Context(), "admit")) - u, h = f.handle(t, a) - rollback := errors.New("rollback archive") var histories []string for _, loaded := range []bool{false, true} { - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - if loaded { - _, err = h.LoadExecution(t.Context()) - require.NoError(t, err) - } - _, err = h.Archive(t.Context(), uuid.Nil) + u, h = f.handle(t, a) + if loaded { + _, err = h.LoadExecution(t.Context()) require.NoError(t, err) - var history string - require.NoError(t, - u.DB().QueryRow(t.Context(), - `SELECT coalesce(jsonb_agg(jsonb_build_object('event',to_jsonb(e)-ARRAY['id','agent_input_id','idempotency_key','created_at'],'input',to_jsonb(i)-ARRAY['id','queued_at','admitted_at','resolved_at','canceled_at','admitted_event_id']) ORDER BY e.sequence),'[]')::text FROM agent_events e LEFT JOIN agent_inputs i ON i.id=e.agent_input_id WHERE e.agent_id=$1`, - a).Scan(&history)) - histories = append(histories, history) - return rollback - }), rollback) + } + _, err = h.Archive(t.Context(), uuid.Nil) + require.NoError(t, err) + var history string + require.NoError(t, + u.DB().QueryRow(t.Context(), + `SELECT coalesce(jsonb_agg(jsonb_build_object('event',to_jsonb(e)-ARRAY['id','agent_input_id','idempotency_key','created_at'],'input',to_jsonb(i)-ARRAY['id','queued_at','admitted_at','resolved_at','canceled_at','admitted_event_id']) ORDER BY e.sequence),'[]')::text FROM agent_events e LEFT JOIN agent_inputs i ON i.id=e.agent_input_id WHERE e.agent_id=$1`, + a).Scan(&history)) + histories = append(histories, history) + require.NoError(t, u.Rollback(t.Context())) } require.Equal(t, histories[0], histories[1]) require.Contains(t, histories[0], `"input_kind": "control"`) diff --git a/internal/storage/internal/agentexecution/commands_permission_integration_test.go b/internal/storage/internal/agentexecution/commands_permission_integration_test.go index af6cc5fde..eb09d93ab 100644 --- a/internal/storage/internal/agentexecution/commands_permission_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_permission_integration_test.go @@ -65,14 +65,16 @@ func TestExecutionPermissionAndPresentation(t *testing.T) { changed, err = h.RecordPresentation(t.Context(), opened.ID, dest, receipt) require.NoError(t, err) require.False(t, changed) - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.RecordPresentation(t.Context(), opened.ID, json.RawMessage(`{}`), receipt) - return err - }), storeerr.ErrUnauthorized) - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.RecordPresentation(t.Context(), opened.ID, dest, json.RawMessage(`{}`)) - return err - }), storeerr.ErrIdempotencyConflict) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.RecordPresentation(t.Context(), opened.ID, json.RawMessage(`{}`), receipt) + require.ErrorIs(t, err, storeerr.ErrUnauthorized) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) + _, err = h.RecordPresentation(t.Context(), opened.ID, dest, json.RawMessage(`{}`)) + require.ErrorIs(t, err, storeerr.ErrIdempotencyConflict) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) snapshot := assertCommandState(t, u, h) if mode == "allow" { require.Equal(t, agentexecution.WorkTool, snapshot.Selection.Work) diff --git a/internal/storage/internal/agentexecution/commands_processes_integration_test.go b/internal/storage/internal/agentexecution/commands_processes_integration_test.go index 00ce4b473..3ba5876e3 100644 --- a/internal/storage/internal/agentexecution/commands_processes_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_processes_integration_test.go @@ -159,10 +159,12 @@ func TestExecutionRejectsMissingDurableOwner(t *testing.T) { ref := toolCommand(t, h, lease, "built_in") _, err := h.AuthorizeTool(t.Context(), ref) require.NoError(t, err) - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.WaitForTool(t.Context(), ref, agentexecution.ToolOwner{ProcessID: uuid.New()}) - return err - }), storeerr.ErrInvalidToolCallDisposition) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.WaitForTool(t.Context(), ref, agentexecution.ToolOwner{ProcessID: uuid.New()}) + require.ErrorIs(t, err, storeerr.ErrInvalidToolCallDisposition) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) snapshot := assertCommandState(t, u, h) require.Equal(t, agentexecution.WorkTool, snapshot.Selection.Work) require.NoError(t, u.Commit(t.Context(), "missing owner rollback")) @@ -229,13 +231,15 @@ func TestExecutionActionReports(t *testing.T) { observed := json.RawMessage( `{"process_id":"` + pid + `","output":"ab","cursor":0,"next_cursor":2,"truncated":false}`, ) - require.ErrorIs(t, u.Savepoint(t.Context(), func(_ *agentexecution.Unit) error { - _, err := h.CompleteProcessAction( - t.Context(), - agentexecution.ActionResult{ID: actions[1], Observed: observed}, - ) - return err - }), storeerr.ErrProcessActionReportBlocked) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.CompleteProcessAction( + t.Context(), + agentexecution.ActionResult{ID: actions[1], Observed: observed}, + ) + require.ErrorIs(t, err, storeerr.ErrProcessActionReportBlocked) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) _, err = u.DB().Exec(t.Context(), `UPDATE process_actions SET state='applied' WHERE id=$1`, actions[0]) require.NoError(t, err) first, err := h.CompleteProcessAction( diff --git a/internal/storage/internal/agentexecution/commands_runtime_integration_test.go b/internal/storage/internal/agentexecution/commands_runtime_integration_test.go index 339e45193..f91b7190c 100644 --- a/internal/storage/internal/agentexecution/commands_runtime_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_runtime_integration_test.go @@ -55,15 +55,11 @@ func TestExecutionAdvanceBudgetsAndRecovery(t *testing.T) { } _, err = u.DB().Exec(t.Context(), query, lease) require.NoError(t, err) - require.ErrorIs( - t, - u.Savepoint( - t.Context(), - func(_ *agentexecution.Unit) error { _, err := h.Advance(t.Context(), lease, true, true); return err }, - ), - storeerr.ErrRuntimeLockInactive, - ) - require.NoError(t, u.Commit(t.Context(), "fenced advance")) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.Advance(t.Context(), lease, true, true) + require.ErrorIs(t, err, storeerr.ErrRuntimeLockInactive) + require.NoError(t, u.Rollback(t.Context())) return } work, err := h.Advance(t.Context(), lease, mode == "model", true) @@ -99,17 +95,12 @@ func TestExecutionClaimRenewAcceptAdvance(t *testing.T) { claimed, err := h.Claim(t.Context(), uuid.New(), 90*time.Second) require.NoError(t, err) require.NotEqual(t, uuid.Nil, claimed.RuntimeID) - require.ErrorIs( - t, - u.Savepoint( - t.Context(), - func(_ *agentexecution.Unit) error { - _, err := h.Claim(t.Context(), uuid.New(), 90*time.Second) - return err - }, - ), - storeerr.ErrAgentNotAdvanceable, - ) + require.NoError(t, u.Commit(t.Context(), "before rejected command")) + u, h = f.handle(t, a) + _, err = h.Claim(t.Context(), uuid.New(), 90*time.Second) + require.ErrorIs(t, err, storeerr.ErrAgentNotAdvanceable) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, a) route, err := h.Route() require.NoError(t, err) _, err = u.RenewRuntime(t.Context(), route, claimed.RuntimeID, 90*time.Second) diff --git a/internal/storage/internal/agentexecution/handle.go b/internal/storage/internal/agentexecution/handle.go index f56e657e0..e9cb79d1e 100644 --- a/internal/storage/internal/agentexecution/handle.go +++ b/internal/storage/internal/agentexecution/handle.go @@ -9,7 +9,6 @@ import ( type Handle struct { unit *Unit route AgentRoute - valid bool created bool mutation *executionMutation identity *dbsqlc.LockAgentInProjectRow @@ -28,7 +27,7 @@ func (u *Unit) Agent(route AgentRoute) (*Handle, error) { } func (h *Handle) Route() (AgentRoute, error) { - if !h.valid || h.unit.handles[h.route.AgentID] != h { + if h.unit == nil || h.unit.handles[h.route.AgentID] != h { return AgentRoute{}, ErrUnitClosed } if err := h.unit.active(); err != nil { diff --git a/internal/storage/internal/agentexecution/locks.go b/internal/storage/internal/agentexecution/locks.go index ce8e3616f..210685627 100644 --- a/internal/storage/internal/agentexecution/locks.go +++ b/internal/storage/internal/agentexecution/locks.go @@ -137,7 +137,7 @@ func (u *Unit) lockAgents(ctx context.Context, plan MutationPlan, skipLocked boo if err != nil { return false, err } - u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true, identity: &identity} + u.handles[route.AgentID] = &Handle{unit: u, route: route, identity: &identity} } if owned := plan.authority; owned != nil { h := u.handles[owned.AgentID] @@ -222,7 +222,7 @@ func (u *Unit) ClaimAgent(ctx context.Context) (AgentRoute, error) { route := AgentRoute{ProjectID: wakeup.ProjectID, AgentID: wakeup.AgentID, RootAgentID: wakeup.RootAgentID} if u.handles[route.AgentID] == nil { - u.handles[route.AgentID] = &Handle{unit: u, route: route, valid: true} + u.handles[route.AgentID] = &Handle{unit: u, route: route} } u.claimedAgent = route.AgentID return route, nil @@ -244,7 +244,7 @@ func (u *Unit) CreatedAgent(projectID, agentID, parentID uuid.UUID) error { rootID = parent.route.RootAgentID } route := AgentRoute{ProjectID: projectID, AgentID: agentID, RootAgentID: rootID} - u.handles[agentID] = &Handle{unit: u, route: route, valid: true, created: true} + u.handles[agentID] = &Handle{unit: u, route: route, created: true} return nil } diff --git a/internal/storage/internal/agentexecution/mutation.go b/internal/storage/internal/agentexecution/mutation.go index d0b08c3f7..377ad24fb 100644 --- a/internal/storage/internal/agentexecution/mutation.go +++ b/internal/storage/internal/agentexecution/mutation.go @@ -19,19 +19,6 @@ type executionMutation struct { loaded *ExecutionSnapshot } -func (m *executionMutation) clone() *executionMutation { - cloned := *m - if m.loaded != nil { - snapshot := cloneExecutionSnapshot(*m.loaded) - cloned.loaded = &snapshot - } - if m.head.LogicalReadyAt != nil { - ready := *m.head.LogicalReadyAt - cloned.head.LogicalReadyAt = &ready - } - return &cloned -} - func (h *Handle) executionMutation(ctx context.Context) (*executionMutation, error) { if _, err := h.Route(); err != nil { return nil, err diff --git a/internal/storage/internal/agentexecution/schema_integration_test.go b/internal/storage/internal/agentexecution/schema_integration_test.go index 997708be0..73fa9eb0b 100644 --- a/internal/storage/internal/agentexecution/schema_integration_test.go +++ b/internal/storage/internal/agentexecution/schema_integration_test.go @@ -3,7 +3,6 @@ package agentexecution_test import ( - "errors" "go/ast" "go/parser" "go/token" @@ -178,19 +177,19 @@ SELECT normal_context_id IS NULL FROM agent_execution_state WHERE agent_id=$1 Scan(&unchanged), ) require.True(t, unchanged) - rejected := errors.New("rollback input") - require.ErrorIs(t, u.Savepoint(t.Context(), func(u *agentexecution.Unit) error { - _, insertErr := u.DB(). - Exec(t.Context(), `INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) - VALUES($1,$2,'received','content','steering',$3)`, f.project, agent, f.now) - require.NoError(t, insertErr) - rebuilt, rebuildErr := h.ReconstructExecution(t.Context()) - require.NoError(t, rebuildErr) - require.Equal(t, agentexecution.AdmitAllSteering, rebuilt.Selection.Admission) - return rejected - }), rejected) + require.NoError(t, u.Commit(t.Context(), "corrupt head")) + u, h = f.handle(t, agent) + _, err = u.DB().Exec(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) + VALUES($1,$2,'received','content','steering',$3)`, f.project, agent, f.now) + require.NoError(t, err) rebuilt, err := h.ReconstructExecution(t.Context()) require.NoError(t, err) + require.Equal(t, agentexecution.AdmitAllSteering, rebuilt.Selection.Admission) + require.NoError(t, u.Rollback(t.Context())) + u, h = f.handle(t, agent) + rebuilt, err = h.ReconstructExecution(t.Context()) + require.NoError(t, err) require.Equal(t, agentexecution.WorkModel, rebuilt.Selection.Work) _, err = u.DB().Exec(t.Context(), `DELETE FROM agent_execution_state WHERE agent_id=$1`, agent) require.NoError(t, err) diff --git a/internal/storage/internal/agentexecution/unit.go b/internal/storage/internal/agentexecution/unit.go index 0ec9017b4..9dd4c6ceb 100644 --- a/internal/storage/internal/agentexecution/unit.go +++ b/internal/storage/internal/agentexecution/unit.go @@ -29,7 +29,6 @@ type Unit struct { closed bool aborted bool finalizing bool - savepoints int } func newUnit(cell *AgentCell, tx pgx.Tx) *Unit { @@ -58,8 +57,8 @@ func (u *Unit) Commit(ctx context.Context, operation string) (err error) { if err := u.active(); err != nil { return err } - if u.savepoints != 0 || u.finalizing { - return errors.New("cannot commit inside unit savepoint or finalization") + if u.finalizing { + return errors.New("cannot commit during unit finalization") } u.finalizing = true defer func() { @@ -87,9 +86,6 @@ func (u *Unit) Rollback(ctx context.Context) error { if u.closed { return pgx.ErrTxClosed } - if u.savepoints != 0 { - return errors.New("use savepoint rollback inside a unit savepoint") - } u.closed = true return u.tx.Rollback(ctx) } @@ -99,87 +95,3 @@ func (u *Unit) orderedHandles() []*Handle { slices.SortFunc(handles, func(a, b *Handle) int { return compareRoutes(a.route, b.route) }) return handles } - -func (u *Unit) Savepoint(ctx context.Context, run func(*Unit) error) (err error) { - if err := u.active(); err != nil { - return err - } - if u.finalizing { - return errors.New("cannot create savepoint during finalization") - } - parent := u.tx - nested, err := parent.Begin(ctx) - if err != nil { - return err - } - snapshot := u.snapshot() - u.tx = nested - u.savepoints++ - released := false - defer func() { - u.savepoints-- - u.tx = parent - if !released { - rollbackErr := nested.Rollback(context.WithoutCancel(ctx)) - u.restore(snapshot) - if rollbackErr != nil { - u.aborted = true - err = errors.Join(err, fmt.Errorf("rollback agent execution savepoint: %w", rollbackErr)) - } - } - }() - if err := run(u); err != nil { - return err - } - if err := u.active(); err != nil { - return err - } - if err := nested.Commit(ctx); err != nil { - return err - } - released = true - return nil -} - -type unitSnapshot struct { - claimedAgent uuid.UUID - handles map[uuid.UUID]*Handle - states map[uuid.UUID]Handle - notifications *notifications.TxNotifications - exclusiveProjects map[uuid.UUID]struct{} - runtimeLocked bool -} - -func (u *Unit) snapshot() unitSnapshot { - s := unitSnapshot{ - claimedAgent: u.claimedAgent, handles: maps.Clone(u.handles), states: make(map[uuid.UUID]Handle), - notifications: u.notifications.Clone(), - exclusiveProjects: maps.Clone(u.exclusiveProjects), - runtimeLocked: u.runtimeLocked, - } - for id, h := range u.handles { - state := *h - if h.mutation != nil { - state.mutation = h.mutation.clone() - } - s.states[id] = state - } - return s -} - -func (u *Unit) restore(snapshot unitSnapshot) { - for id, h := range u.handles { - if snapshot.handles[id] != h { - h.valid = false - } - } - u.handles = snapshot.handles - for id, h := range u.handles { - *h = snapshot.states[id] - } - u.notifications.Restore(snapshot.notifications) - u.exclusiveProjects = snapshot.exclusiveProjects - u.runtimeLocked = snapshot.runtimeLocked - u.claimedAgent = snapshot.claimedAgent - u.aborted = false -} diff --git a/internal/storage/internal/agentexecution/unit_integration_test.go b/internal/storage/internal/agentexecution/unit_integration_test.go index 1cae41bd1..de553b959 100644 --- a/internal/storage/internal/agentexecution/unit_integration_test.go +++ b/internal/storage/internal/agentexecution/unit_integration_test.go @@ -4,7 +4,6 @@ package agentexecution import ( "context" - "errors" "testing" "github.com/google/uuid" @@ -95,76 +94,6 @@ func TestUnitCommitLifecycle(t *testing.T) { require.ErrorIs(t, u.Commit(ctx, "again"), ErrUnitClosed) } -func TestUnitSavepointSnapshots(t *testing.T) { - ctx := t.Context() - var published []notifications.PostCommitIntent - cell, pool := unitTestCell( - t, - unitPublisher{publish: func(_ context.Context, intent notifications.PostCommitIntent) { - published = append(published, intent) - }}, - ) - u, err := BeginTestUnit(ctx, cell) - require.NoError(t, err) - defer func() { _ = u.Rollback(ctx) }() - project, agent, machine := uuid.New(), uuid.New(), uuid.New() - h := unitTestAgent(t, u, project, agent) - mutation := &executionMutation{head: ExecutionHead{StopSequence: 1}} - h.mutation = mutation - pending := u.Notifications() - pending.AddDaemonWork(machine) - process := uuid.New() - pending.AddDaemonProcessTermination(machine, process) - rejected := errors.New("reject inner") - var escaped *Handle - require.ErrorIs(t, u.Savepoint(ctx, func(u *Unit) error { - mutation.head.StopSequence = 2 - pending.AddDaemonWork(uuid.New()) - pending.AddDaemonProcessTermination(machine, uuid.New()) - delete(u.handles, agent) - escaped = unitTestAgent(t, u, project, uuid.New()) - escaped.mutation = &executionMutation{head: ExecutionHead{StopSequence: 3}} - delete(u.handles, escaped.route.AgentID) - _, err := u.DB(). - Exec(ctx, `INSERT INTO unit_probe(phase) VALUES ('rollback'); SELECT pg_advisory_xact_lock(8173)`) - require.NoError(t, err) - require.NoError(t, u.Savepoint(ctx, func(u *Unit) error { - _, err := u.DB().Exec(ctx, `INSERT INTO unit_probe(phase) VALUES ('nested release')`) - return err - })) - require.Empty(t, published) - return rejected - }), rejected) - require.Same(t, pending, u.Notifications()) - require.Same(t, h, u.handles[agent]) - restored := h.mutation - require.EqualValues(t, 1, restored.head.StopSequence) - _, err = escaped.Route() - require.ErrorIs(t, err, ErrUnitClosed) - var unlocked bool - require.NoError(t, pool.QueryRow(ctx, `SELECT pg_try_advisory_xact_lock(8173)`).Scan(&unlocked)) - require.True(t, unlocked) - require.NoError(t, u.Savepoint(ctx, func(u *Unit) error { - require.Error(t, u.Commit(ctx, "nested")) - require.Error(t, u.Rollback(ctx)) - restored.head.StopSequence = 4 - return nil - })) - require.NoError(t, u.Commit(ctx, "snapshot")) - require.Len(t, published, 2) - for _, intent := range published { - if terminate, ok := intent.(notifications.DaemonProcessTerminationCommitted); ok { - require.Equal(t, []uuid.UUID{process}, terminate.ProcessIDs) - } - } - var phases []string - require.NoError( - t, - pool.QueryRow(ctx, `SELECT array_agg(phase ORDER BY id) FROM unit_probe`).Scan(&phases), - ) - require.Empty(t, phases) -} - func TestUnitFailuresSuppressPublication(t *testing.T) { for _, scenario := range []string{"rollback", "deferred commit", "aborted SQL", "cancel"} { t.Run(scenario, func(t *testing.T) { @@ -205,23 +134,7 @@ func TestUnitFailuresSuppressPublication(t *testing.T) { } } -func TestUnitSavepointRecoversDatabaseErrors(t *testing.T) { - ctx := t.Context() - cell, _ := unitTestCell(t, nil) - require.NoError(t, cell.Transact(ctx, func(u *Unit) error { - err := u.Savepoint(ctx, func(u *Unit) error { - _, err := u.DB().Exec(ctx, `SELECT 1 / 0`) - require.Error(t, err) - return nil - }) - require.ErrorIs(t, err, ErrUnitAborted) - _, err = u.DB().Exec(ctx, `INSERT INTO deferred_probe VALUES (1)`) - require.NoError(t, err) - return err - })) -} - -func TestUnitLockPlansAndSavepointLocks(t *testing.T) { +func TestUnitLockPlans(t *testing.T) { ctx := t.Context() cell, pool := unitTestCell(t, nil) project, org := uuid.New(), uuid.New() @@ -255,18 +168,7 @@ func TestUnitLockPlansAndSavepointLocks(t *testing.T) { require.NoError(t, err) defer func() { _ = other.Rollback(ctx) }() require.ErrorIs(t, other.LockAgents(ctx, plan), ErrLockPlan) - rollback := errors.New("rollback locks") - require.ErrorIs(t, other.Savepoint(ctx, func(u *Unit) error { - return errors.Join(u.LockAgentRefs( - ctx, - []lifecyclelock.AgentRef{{ProjectID: project, AgentID: lower}}, nil), - rollback, - ) - }), rollback) - require.Empty(t, other.handles) - check := integrationdb.BeginTx(t, ctx, pool) - _, err = check.Exec(ctx, `SELECT id FROM agents WHERE id=$1 FOR UPDATE NOWAIT`, lower) - require.NoError(t, err) + } func TestUnitExpandsLockPlanAndRollsBackAttempt(t *testing.T) { @@ -417,46 +319,3 @@ func TestCellBeginRequiresRetryBoundary(t *testing.T) { } require.ErrorContains(t, err, "retry boundary") } - -func TestUnitSavepointPropagatesLockPlan(t *testing.T) { - ctx := t.Context() - cell, pool := unitTestCell(t, nil) - project, org, first, second := uuid.New(), uuid.New(), uuid.New(), uuid.New() - _, err := pool.Exec(ctx, `INSERT INTO agents(id,project_id,org_id,root_agent_id) -SELECT id,$1,$2,id FROM unnest($3::uuid[]) id`, project, org, []uuid.UUID{first, second}) - require.NoError(t, err) - attempts := 0 - _, err = Retry(ctx, "savepoint expansion", func(ctx context.Context) (struct{}, error) { - attempts++ - unit, err := BeginTestUnit(ctx, cell) - if err != nil { - return struct{}{}, err - } - defer func() { _ = unit.Rollback(ctx) }() - if err = unit.LockAgentRefs(ctx, []lifecyclelock.AgentRef{{ProjectID: project, AgentID: first}}, nil); err != nil { - return struct{}{}, err - } - err = unit.Savepoint(ctx, func(u *Unit) error { - _, err := u.DB().Exec(ctx, `INSERT INTO unit_probe(agent_id,phase) VALUES($1,'advance')`, first) - if err != nil { - return err - } - _, err = u.Handle(project, second) - return err - }) - if attempts == 1 { - var missing *LockPlanError - require.ErrorAs(t, err, &missing) - require.Contains(t, missing.Agents, lifecyclelock.AgentRef{ProjectID: project, AgentID: second}) - } - if err != nil { - return struct{}{}, err - } - return struct{}{}, unit.Commit(ctx, "expanded continuation") - }) - require.NoError(t, err) - require.Equal(t, 2, attempts) - var count int - require.NoError(t, pool.QueryRow(ctx, `SELECT count(*) FROM unit_probe WHERE agent_id=$1`, first).Scan(&count)) - require.Equal(t, 1, count) -} diff --git a/internal/storage/internal/dbsqlc/agent_inputs.sql.go b/internal/storage/internal/dbsqlc/agent_inputs.sql.go index 050cc5fc5..17d398d1e 100644 --- a/internal/storage/internal/dbsqlc/agent_inputs.sql.go +++ b/internal/storage/internal/dbsqlc/agent_inputs.sql.go @@ -20,11 +20,6 @@ WITH locked_agent AS MATERIALIZED ( JOIN agents agent ON agent.id = wake.agent_id WHERE agent.state <> 'archived' AND wake.ready_at <= statement_timestamp() - AND NOT EXISTS ( - SELECT 1 - FROM agent_runtime_locks runtime_lock - WHERE runtime_lock.agent_id = wake.agent_id - ) ORDER BY wake.ready_at ASC, wake.agent_id ASC FOR UPDATE OF agent SKIP LOCKED LIMIT 1 @@ -50,10 +45,6 @@ type ClaimNextAgentWakeupRow struct { RootAgentID uuid.UUID } -// The query walks wakeups in global ready order, but locks agents before wake -// rows to match all other agent mutation paths. If the selected wake is no -// longer claimable after the runtime-lock recheck, this returns no rows and the -// worker retries on its next poll; no work is lost. func (q *Queries) ClaimNextAgentWakeup(ctx context.Context) (ClaimNextAgentWakeupRow, error) { row := q.db.QueryRow(ctx, claimNextAgentWakeup) var i ClaimNextAgentWakeupRow diff --git a/internal/storage/queries/agent_inputs.sql b/internal/storage/queries/agent_inputs.sql index 50a2dbfc9..f1c9125f7 100644 --- a/internal/storage/queries/agent_inputs.sql +++ b/internal/storage/queries/agent_inputs.sql @@ -1,19 +1,10 @@ -- name: ClaimNextAgentWakeup :one --- The query walks wakeups in global ready order, but locks agents before wake --- rows to match all other agent mutation paths. If the selected wake is no --- longer claimable after the runtime-lock recheck, this returns no rows and the --- worker retries on its next poll; no work is lost. WITH locked_agent AS MATERIALIZED ( SELECT agent.id AS agent_id, agent.project_id, agent.root_agent_id FROM agent_wakeups wake JOIN agents agent ON agent.id = wake.agent_id WHERE agent.state <> 'archived' AND wake.ready_at <= statement_timestamp() - AND NOT EXISTS ( - SELECT 1 - FROM agent_runtime_locks runtime_lock - WHERE runtime_lock.agent_id = wake.agent_id - ) ORDER BY wake.ready_at ASC, wake.agent_id ASC FOR UPDATE OF agent SKIP LOCKED LIMIT 1 From e2ab1eff85ad7d7e72ead5ff0d607593acbd93a0 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Mon, 12 Oct 2026 05:57:09 -0700 Subject: [PATCH 06/19] test: pin MultiXact-free fused acceptance and the ready-index claim plan Co-Authored-By: Claude Opus 5.5 (1M context) --- .../model_work_multixact_integration_test.go | 86 +++++++++++++++++++ .../claim_plan_integration_test.go | 61 +++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 internal/storage/executionstore/model_work_multixact_integration_test.go create mode 100644 internal/storage/internal/agentexecution/claim_plan_integration_test.go diff --git a/internal/storage/executionstore/model_work_multixact_integration_test.go b/internal/storage/executionstore/model_work_multixact_integration_test.go new file mode 100644 index 000000000..ed9a9abaf --- /dev/null +++ b/internal/storage/executionstore/model_work_multixact_integration_test.go @@ -0,0 +1,86 @@ +//go:build integration + +package executionstore_test + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/omnara-ai/omnara/internal/modelenvelope" + "github.com/stretchr/testify/require" +) + +func TestFusedAcceptanceDoesNotCreateMultiXacts(t *testing.T) { + for _, continued := range []bool{true, false} { + name := "release" + if continued { + name = "continue" + } + t.Run(name, func(t *testing.T) { + ctx := t.Context() + f, input := fusedModelFixture(t, ctx) + if continued { + input.Output.ProviderResponse.Normalized.StopReason = modelenvelope.StopReasonMaxTokens + } + _, err := f.Store.pool.Exec(ctx, `CREATE EXTENSION pageinspect; +ALTER TABLE agents SET (autovacuum_enabled=false); +ALTER TABLE agent_runtime_locks SET (autovacuum_enabled=false)`) + require.NoError(t, err) + multixacts := func(table string) int { + t.Helper() + var n int + require.NoError(t, f.Store.pool.QueryRow(ctx, ` +SELECT count(*) FROM generate_series(0,pg_relation_size($1::text::regclass)/current_setting('block_size')::int-1) block +CROSS JOIN LATERAL heap_page_items(get_raw_page($1,block)) item +CROSS JOIN LATERAL heap_tuple_infomask_flags(item.t_infomask,item.t_infomask2) flags +WHERE flags.raw_flags @> ARRAY['HEAP_XMAX_IS_MULTI']`, table).Scan(&n)) + return n + } + for _, table := range []string{"agents", "agent_runtime_locks"} { + require.Zero(t, multixacts(table), table+" before") + } + counter := &fusedQueryCounter{} + config := f.Store.pool.Config() + config.ConnConfig.Tracer = counter + pool, err := pgxpool.NewWithConfig(ctx, config) + require.NoError(t, err) + defer pool.Close() + measured := newIntegrationStore(pool) + transition, err := measured.Execution().AcceptModelOutputAndAdvance(ctx, input) + t.Logf("fused %s statements=%d", name, counter.statements) + bound := 20 + if continued { + bound = 29 + } + require.LessOrEqual(t, counter.statements, bound) + require.NoError(t, err) + require.Nil(t, transition.Fallback) + require.Equal(t, continued, transition.Continued) + for _, table := range []string{"agents", "agent_runtime_locks"} { + n := multixacts(table) + t.Logf("%s MultiXact tuples after fused acceptance: %d", table, n) + require.Zero(t, n, table) + } + }) + } +} + +type fusedQueryCounter struct{ statements int } + +func (c *fusedQueryCounter) TraceQueryStart( + ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData, +) context.Context { + c.statements++ + return ctx +} +func (*fusedQueryCounter) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {} +func (*fusedQueryCounter) TraceBatchStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceBatchStartData) context.Context { + return ctx +} +func (c *fusedQueryCounter) TraceBatchQuery(context.Context, *pgx.Conn, pgx.TraceBatchQueryData) { + c.statements++ +} +func (*fusedQueryCounter) TraceBatchEnd(context.Context, *pgx.Conn, pgx.TraceBatchEndData) {} diff --git a/internal/storage/internal/agentexecution/claim_plan_integration_test.go b/internal/storage/internal/agentexecution/claim_plan_integration_test.go new file mode 100644 index 000000000..803e95630 --- /dev/null +++ b/internal/storage/internal/agentexecution/claim_plan_integration_test.go @@ -0,0 +1,61 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "os" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestClaimReadyIndexPlan(t *testing.T) { + f := migratedExecutionFixture(t) + ctx := t.Context() + ids := make([]uuid.UUID, 10000) + for i := range ids { + ids[i] = uuid.New() + } + _, err := f.pool.Exec(ctx, ` +INSERT INTO agents(id,root_agent_id,org_id,project_id,state,current_config_id,created_at,updated_at) + SELECT id,id,$2,$3,'active',$4,statement_timestamp(),statement_timestamp() FROM unnest($1::uuid[]) id`, + ids, f.org, f.project, f.config) + require.NoError(t, err) + _, err = f.pool.Exec(ctx, `INSERT INTO agent_wakeups(agent_id,ready_at,updated_at) + SELECT id,statement_timestamp()-interval '1 minute',statement_timestamp() FROM unnest($1::uuid[]) id`, ids[:5000]) + require.NoError(t, err) + _, err = f.pool.Exec(ctx, ` +INSERT INTO agent_runtime_locks(id,agent_id,worker_process_id,started_at,renewed_at,lease_expires_at) + SELECT gen_random_uuid(),id,$2,statement_timestamp(),statement_timestamp(),statement_timestamp()+interval '90 seconds' + FROM unnest($1::uuid[]) id`, ids[5000:], uuid.New()) + require.NoError(t, err) + _, err = f.pool.Exec(ctx, `ANALYZE agents; ANALYZE agent_wakeups; ANALYZE agent_runtime_locks`) + require.NoError(t, err) + sql, err := os.ReadFile("../../queries/agent_inputs.sql") + require.NoError(t, err) + query := strings.Split(string(sql), "-- name: LockClaimableAgentWakeup")[0] + tx, err := f.pool.Begin(ctx) + require.NoError(t, err) + defer func() { _ = tx.Rollback(ctx) }() + _, err = tx.Exec(ctx, `SET LOCAL plan_cache_mode=force_generic_plan`) + require.NoError(t, err) + var raw []byte + require.NoError(t, tx.QueryRow(ctx, "EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) "+query).Scan(&raw)) + t.Logf("claim EXPLAIN: %s", raw) + var plan []struct { + Plan struct { + SharedHitBlocks int `json:"Shared Hit Blocks"` + SharedReadBlocks int `json:"Shared Read Blocks"` + } + } + require.NoError(t, json.Unmarshal(raw, &plan)) + require.NotContains(t, string(raw), `"Node Type": "Sort"`) + require.NotContains(t, string(raw), `"Node Type": "Hash"`) + require.NotContains(t, string(raw), `"Node Type": "Hash Join"`) + require.Contains(t, string(raw), "agent_wakeups_global_ready_idx") + buffers := plan[0].Plan.SharedHitBlocks + plan[0].Plan.SharedReadBlocks + require.LessOrEqual(t, buffers, 40) +} From 3ad8a4514286997769a13f14f056d40ab9bd17d8 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:35:30 -0700 Subject: [PATCH 07/19] perf: index lease references and scope hot inserts to the agent row Fold into the 052 cutover: - Index tool_calls.runtime_lock_id so the foreign-key check on every lease delete no longer scans all tool-call history. - Index received content inputs per agent for the execution backlog. - Replace the org and project references on model contexts, processes and process actions with one org-qualified reference to the agent row, which the writer already holds locked, and drop the redundant project reference on agent inputs. Inserts stop key-share locking the shared org and project rows. Config and model-revision references stay. - Move the payload-time helper from 053 into 052, so the cutover is a single migration. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../execution_cutover_integration_test.go | 103 +++++++- .../hot_path_plans_integration_test.go | 242 ++++++++++++++++++ .../hot_path_schema_integration_test.go | 230 +++++++++++++++++ migrations/000052_agent_execution_state.sql | 66 +++++ .../000053_agent_payload_query_plans.sql | 30 --- 5 files changed, 638 insertions(+), 33 deletions(-) create mode 100644 internal/storage/internal/agentexecution/hot_path_plans_integration_test.go create mode 100644 internal/storage/internal/agentexecution/hot_path_schema_integration_test.go delete mode 100644 migrations/000053_agent_payload_query_plans.sql diff --git a/internal/dbmigrate/execution_cutover_integration_test.go b/internal/dbmigrate/execution_cutover_integration_test.go index da2f83888..3f4bc0415 100644 --- a/internal/dbmigrate/execution_cutover_integration_test.go +++ b/internal/dbmigrate/execution_cutover_integration_test.go @@ -3,17 +3,61 @@ package dbmigrate_test import ( - "context" + "os" "strings" "testing" + "testing/fstest" + "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/pgx/v5/stdlib" + "github.com/omnara-ai/omnara/internal/dbmigrate" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" "github.com/stretchr/testify/require" ) func TestExecutionCutoverRemovesHistoryEngine(t *testing.T) { t.Parallel() - ctx := context.Background() - pool, _ := openPostgresMigrationTestDB(t, ctx) + for _, mode := range []string{"fresh", "released", "rollback"} { + t.Run(mode, func(t *testing.T) { + t.Parallel() + ctx := t.Context() + if mode == "fresh" { + pool, _ := openPostgresMigrationTestDB(t, ctx) + assertExecutionCutoverCatalog(t, pool) + return + } + pool := integrationdb.OpenUnmigratedPool(t, ctx) + db := stdlib.OpenDBFromPool(pool) + defer func() { _ = db.Close() }() + require.NoError(t, applyProductionPostgresMigrationsThrough(t, ctx, db, 51)) + if mode == "rollback" { + body, err := os.ReadFile("../../migrations/000052_agent_execution_state.sql") + require.NoError(t, err) + body = append(body, []byte("\nSELECT gate_11_cutover_failure();\n")...) + files := fstest.MapFS{"000052_agent_execution_state.sql": {Data: body}} + require.ErrorContains(t, dbmigrate.ApplyPostgres(ctx, db, files), "gate_11_cutover_failure") + require.EqualValues(t, 51, currentPostgresMigrationVersion(t, ctx, db)) + var restored bool + require.NoError(t, pool.QueryRow(ctx, ` +SELECT to_regclass('tool_calls_runtime_lock_idx') IS NOT NULL + AND to_regclass('tool_calls_runtime_reference_idx') IS NULL + AND to_regclass('agent_inputs_received_content_idx') IS NULL + AND to_regclass('agents_org_project_id_key') IS NULL + AND to_regclass('agent_execution_state') IS NULL + AND to_regprocedure('agent_payload_created_at(uuid,uuid,text,uuid,boolean)') IS NULL + AND (SELECT count(*) FROM pg_constraint WHERE conname=ANY($1::text[]))=9 +`, replacedExecutionScopeConstraints()).Scan(&restored)) + require.True(t, restored) + } + require.NoError(t, applyProductionPostgresMigrations(ctx, db)) + assertExecutionCutoverCatalog(t, pool) + }) + } +} + +func assertExecutionCutoverCatalog(t *testing.T, pool *pgxpool.Pool) { + t.Helper() + ctx := t.Context() names := []string{ "agent_scheduler_state", "model_call_context_turns", @@ -67,4 +111,57 @@ func TestExecutionCutoverRemovesHistoryEngine(t *testing.T) { } } require.NoError(t, rows.Err()) + for _, table := range []string{"model_call_contexts", "processes", "process_actions"} { + var definition string + var validated bool + require.NoError(t, pool.QueryRow(ctx, ` +SELECT pg_get_constraintdef(oid),convalidated FROM pg_constraint + WHERE conrelid=$1::regclass AND conname=$2 +`, table, table+"_agent_scope_fkey").Scan(&definition, &validated)) + require.True(t, validated) + require.Equal(t, "FOREIGN KEY (org_id, project_id, agent_id) REFERENCES agents(org_id, project_id, id)", definition) + } + var replaced int + require.NoError(t, pool.QueryRow(ctx, ` +SELECT count(*) FROM pg_constraint WHERE conname=ANY($1::text[]) +`, replacedExecutionScopeConstraints()).Scan(&replaced)) + require.Zero(t, replaced) + for _, index := range []struct { + name, definition string + unique bool + }{ + {"tool_calls_runtime_reference_idx", "USING btree (runtime_lock_id) WHERE (runtime_lock_id IS NOT NULL)", false}, + {"agent_inputs_received_content_idx", "USING btree (agent_id, input_rank, queued_at, id)", false}, + {"agents_org_project_id_key", "USING btree (org_id, project_id, id)", true}, + } { + var definition string + var ready, valid, unique bool + require.NoError(t, pool.QueryRow(ctx, ` +SELECT pg_get_indexdef(indexrelid),indisready,indisvalid,indisunique FROM pg_index WHERE indexrelid=$1::regclass +`, index.name).Scan(&definition, &ready, &valid, &unique)) + require.True(t, ready) + require.True(t, valid) + require.Equal(t, index.unique, unique) + require.Contains(t, definition, index.definition) + if index.name == "agent_inputs_received_content_idx" { + require.Contains(t, definition, "state = 'received'::text") + require.Contains(t, definition, "input_kind = 'content'::text") + } + } + var payloadInstalled, oldIndexAbsent bool + require.NoError(t, pool.QueryRow(ctx, ` +SELECT to_regprocedure('agent_payload_created_at(uuid,uuid,text,uuid,boolean)') IS NOT NULL, + to_regclass('tool_calls_runtime_lock_idx') IS NULL +`).Scan(&payloadInstalled, &oldIndexAbsent)) + require.True(t, payloadInstalled) + require.True(t, oldIndexAbsent) +} + +func replacedExecutionScopeConstraints() []string { + return []string{ + "model_call_contexts_org_id_fkey", "model_call_contexts_org_id_project_id_fkey", + "model_call_contexts_project_id_agent_id_fkey", "agent_inputs_project_id_fkey", + "processes_org_id_fkey", "processes_org_id_project_id_fkey", "processes_project_id_agent_id_fkey", + "process_actions_org_id_fkey", "process_actions_org_id_project_id_fkey", + } } diff --git a/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go b/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go new file mode 100644 index 000000000..0f0b5208a --- /dev/null +++ b/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go @@ -0,0 +1,242 @@ +//go:build integration + +package agentexecution_test + +import ( + "go/ast" + "go/parser" + "go/token" + "strconv" + "strings" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func generatedQuery(t *testing.T, path, name string) string { + t.Helper() + file, err := parser.ParseFile(token.NewFileSet(), path, nil, 0) + require.NoError(t, err) + var query string + ast.Inspect(file, func(node ast.Node) bool { + literal, ok := node.(*ast.BasicLit) + if !ok || literal.Kind != token.STRING { + return true + } + text, err := strconv.Unquote(literal.Value) + require.NoError(t, err) + if strings.HasPrefix(text, "-- name: "+name+" ") { + query = text + } + return true + }) + require.NotEmpty(t, query, name) + return query +} + +func TestExecutionHotPathPlans(t *testing.T) { + f, agent, lease := commandFixture(t) + ids := make([]uuid.UUID, 2000) + for i := range ids { + ids[i] = uuid.New() + } + _, err := f.pool.Exec(t.Context(), ` +INSERT INTO agents(id,root_agent_id,org_id,project_id,state,current_config_id,created_at,updated_at) + SELECT id,id,$2,$3,'active',$4,now(),now() FROM unnest($1::uuid[]) id +`, ids, f.org, f.project, f.config) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) + SELECT $2,id,'received','content',CASE WHEN n%2=0 THEN 'steering' ELSE 'queued' END,now() + FROM unnest($1::uuid[]) id CROSS JOIN generate_series(1,10) n +`, ids, f.project) + require.NoError(t, err) + seed, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = seed.Rollback(t.Context()) }() + _, err = seed.Exec(t.Context(), ` +INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,ordinal,block_kind,text_content, +created_at) +SELECT agent_id,'agent_input',id,0,'text','history',queued_at FROM agent_inputs WHERE +agent_id=ANY($1::uuid[]) +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +INSERT INTO agent_events(agent_id,turn_id,sequence,event_kind,agent_input_id,is_opening_event, +created_at,idempotency_key) + SELECT agent_id,agent_id,row_number() OVER(PARTITION BY agent_id ORDER BY id),'agent_input',id, +row_number() OVER(PARTITION BY agent_id ORDER BY id)=1,queued_at,'agent_input:'||id FROM agent_inputs +WHERE agent_id=ANY($1::uuid[]) +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +INSERT INTO agent_turns(id,agent_id,turn_sequence,latest_event_id,latest_semantic_event_id) + SELECT agent_id,agent_id,1,id,id FROM agent_events WHERE agent_id=ANY($1::uuid[]) AND sequence=10 +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +UPDATE agent_inputs i SET state='resolved',admitted_event_id=e.id,admitted_at=e.created_at, +resolved_at=e.created_at +FROM agent_events e WHERE i.agent_id=ANY($1::uuid[]) AND e.agent_id=i.agent_id AND +e.agent_input_id=i.id +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), `ANALYZE agent_inputs; ANALYZE agent_events; ANALYZE agent_turns`) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,turn_id,operation_kind,attempt_number, +agent_config_id,configured_model_revision_id,input_event_sequence,runtime_lock_id,created_at, +opening_input_ids,opening_event_sequence) +SELECT e.agent_id,$2,$3,e.agent_id,e.turn_id,'normal',1,$4,$5,10,gen_random_uuid(),now(), +ARRAY[e.agent_input_id],1 + FROM agent_events e WHERE e.agent_id=ANY($1::uuid[]) AND e.sequence=1 +`, ids, f.org, f.project, f.config, f.revision) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +INSERT INTO model_outputs(id,agent_id,model_call_context_id,stop_reason,created_at) + SELECT id,id,id,'end_turn',now() FROM unnest($1::uuid[]) id +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,text_content, +created_at) + SELECT id,'model_output',id,0,'text','answer',now() FROM unnest($1::uuid[]) id +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +INSERT INTO agent_events(agent_id,turn_id,sequence,event_kind,model_output_id,created_at) + SELECT id,id,11,'model_output',id,now() FROM unnest($1::uuid[]) id +`, ids) + require.NoError(t, err) + _, err = seed.Exec(t.Context(), ` +UPDATE model_call_contexts SET state='succeeded',api_format='openai-responses',api_variant='default', +completed_at=now() + WHERE agent_id=ANY($1::uuid[]) +`, ids) + require.NoError(t, err) + require.NoError(t, seed.Commit(t.Context())) + _, err = f.pool.Exec(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) + SELECT $1,$2,'received','content','queued',now() FROM generate_series(1,10000) +`, f.project, agent) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), ` +UPDATE agent_inputs SET state='canceled',canceled_at=now() WHERE agent_id=$1 +`, agent) + require.NoError(t, err) + u, h := f.handle(t, agent) + for range 100 { + receiveCommand(t, h, "steering", "transcript") + _, err = h.AdmitInputs(t.Context()) + require.NoError(t, err) + } + snapshot, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "history fixture")) + var input uuid.UUID + require.NoError(t, f.pool.QueryRow(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) VALUES($1,$2, +'received','content','queued',now()) RETURNING id +`, f.project, agent).Scan(&input)) + _, err = f.pool.Exec(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) + SELECT $2,id,'received','content','queued',now() FROM unnest($1::uuid[]) id +`, ids, f.project) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), ` +ANALYZE agents; ANALYZE agent_inputs; ANALYZE agent_events; ANALYZE content_blocks; ANALYZE +model_call_contexts; ANALYZE model_outputs +`) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), `VACUUM ANALYZE agent_inputs`) + require.NoError(t, err) + tests := []struct { + name, path string + args []any + }{ + {"ListExecutionBacklog", "internal/executiondb/commands_inputs.sql.go", []any{agent}}, + {"ListContextEvents", "../dbsqlc/agent_turns.sql.go", []any{f.project, agent, int64(0), int64(100), int32(100)}}, + {"CreateExecutionAttempt", "internal/executiondb/commands_models.sql.go", []any{uuid.New(), + snapshot.Head.CurrentTurnID, "normal", int32(1), f.config, f.revision, int64(100), nil, lease, + snapshot.Selection.Model.Opening.InputIDs, snapshot.Selection.Model.Opening.EventSequence, agent, + f.project}}, + {"AppendExecutionEvent", "internal/executiondb/commands_events.sql.go", []any{uuid.New(), agent, + snapshot.Head.CurrentTurnID, "agent_input", input, nil, nil, nil, false, nil, f.project}}, + {"InsertExecutionContent", "internal/executiondb/commands_events.sql.go", []any{agent, + "agent_input", input, nil, nil, int32(0), "text", "input", nil, nil, nil, false, []byte(`{}`), f.project}}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + _, err = tx.Exec(t.Context(), `SET LOCAL plan_cache_mode=force_generic_plan`) + require.NoError(t, err) + query := generatedQuery(t, tc.path, tc.name) + var raw []byte + for iteration := range 2 { + nested, err := tx.Begin(t.Context()) + require.NoError(t, err) + require.NoError(t, nested.QueryRow(t.Context(), "EXPLAIN (ANALYZE,BUFFERS,FORMAT JSON) "+query, + tc.args...).Scan(&raw)) + t.Logf("%s iteration=%d plan=%s", tc.name, iteration, raw) + require.NoError(t, nested.Rollback(t.Context())) + } + if tc.name == "ListExecutionBacklog" { + results, err := tx.Conn().PgConn().Exec(t.Context(), "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+query).ReadAll() + require.NoError(t, err) + require.Len(t, results, 1) + require.Contains(t, string(results[0].Rows[0][0]), "agent_inputs_received_content_idx") + require.LessOrEqual(t, planBuffers(t, raw), 10) + require.Contains(t, string(raw), "agent_inputs_received_content_idx") + } + }) + } +} + +func TestContextInsertAvoidsSharedScopeLocks(t *testing.T) { + f := migratedExecutionFixture(t) + _, err := f.pool.Exec(t.Context(), `CREATE EXTENSION pageinspect`) + require.NoError(t, err) + query := ` +INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,turn_id,operation_kind,attempt_number, +agent_config_id,configured_model_revision_id,input_event_sequence,runtime_lock_id,created_at, +opening_input_ids,opening_event_sequence) + VALUES($1,$2,$3,$4,$5,'normal',1,$6,$7,1,$8,now(),ARRAY[$9::uuid],1) +` + type opening struct{ agent, turn, input uuid.UUID } + var openings []opening + for range 2 { + agent := f.create(t, uuid.New(), uuid.Nil) + turn, input := f.opening(t, agent, 1, uuid.Nil, "content") + openings = append(openings, opening{agent, turn, input}) + } + for _, o := range openings { + agent, turn, input := o.agent, o.turn, o.input + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + t.Cleanup(func() { _ = tx.Rollback(t.Context()) }) + _, err = tx.Exec(t.Context(), `SELECT id FROM agents WHERE id=$1 FOR UPDATE`, agent) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), query, uuid.New(), f.org, f.project, agent, turn, f.config, + f.revision, uuid.New(), input) + require.NoError(t, err) + } + for _, table := range []string{"orgs", "projects", "agent_configs"} { + var n int + require.NoError(t, f.pool.QueryRow(t.Context(), ` +SELECT count(*) FROM generate_series(0, +pg_relation_size($1::text::regclass)/current_setting('block_size')::int-1) block + CROSS JOIN LATERAL heap_page_items(get_raw_page($1,block)) item + CROSS JOIN LATERAL heap_tuple_infomask_flags(item.t_infomask,item.t_infomask2) flags + WHERE flags.raw_flags @> ARRAY['HEAP_XMAX_IS_MULTI'] +`, table).Scan(&n)) + t.Logf("%s MultiXact tuples=%d", table, n) + if table == "agent_configs" { + require.Positive(t, n, table) + } else { + require.Zero(t, n, table) + } + } +} diff --git a/internal/storage/internal/agentexecution/hot_path_schema_integration_test.go b/internal/storage/internal/agentexecution/hot_path_schema_integration_test.go new file mode 100644 index 000000000..81a1abff5 --- /dev/null +++ b/internal/storage/internal/agentexecution/hot_path_schema_integration_test.go @@ -0,0 +1,230 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestExecutionLeaseDeleteHistoryBound(t *testing.T) { + f, agent, lease := commandFixture(t) + u, h := f.handle(t, agent) + receiveCommand(t, h, "queued", "input") + _, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + prepared := prepareCommand(t, h, lease) + require.NoError(t, u.Commit(t.Context(), "prepare fixture")) + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + output := uuid.New() + _, err = tx.Exec(t.Context(), ` +INSERT INTO model_outputs(id,agent_id,model_call_context_id,stop_reason,created_at) VALUES($1,$2,$3, +'tool_use',now()) +`, output, agent, prepared.Context.ID) + require.NoError(t, err) + previous := 0 + for _, size := range []int{1000, 100000} { + _, err = tx.Exec(t.Context(), ` +INSERT INTO tool_calls(agent_id,model_output_id,provider_call_id,name,input,type,state,created_at) +SELECT $1,$2,'history-'||n,'tool','{}','built_in','awaiting_authorization',now() FROM +generate_series($3::int,$4::int) n +`, agent, output, previous+1, size) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,tool_call_id, +created_at) +SELECT agent_id,'model_output',model_output_id,split_part(provider_call_id,'-',2)::int,'tool_call',id, +created_at FROM tool_calls WHERE agent_id=$1 AND state='awaiting_authorization' +`, agent) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +UPDATE tool_calls SET state='completed' WHERE agent_id=$1 AND state='awaiting_authorization' +`, agent) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +INSERT INTO tool_call_results(agent_id,tool_call_id,outcome,completed_at) +SELECT agent_id,id,'succeeded',now() FROM tool_calls WHERE agent_id=$1 AND NOT EXISTS(SELECT 1 FROM +tool_call_results r WHERE r.agent_id=$1 AND r.tool_call_id=tool_calls.id) +`, agent) + require.NoError(t, err) + previous = size + _, err = tx.Exec(t.Context(), `ANALYZE tool_calls; SET LOCAL plan_cache_mode=force_generic_plan`) + require.NoError(t, err) + var raw []byte + require.NoError(t, tx.QueryRow(t.Context(), ` +EXPLAIN (ANALYZE,BUFFERS,FORMAT JSON) SELECT 1 FROM ONLY tool_calls x WHERE runtime_lock_id=$1 FOR KEY +SHARE OF x +`, lease).Scan(&raw)) + results, err := tx.Conn().PgConn().Exec(t.Context(), ` +EXPLAIN (GENERIC_PLAN,FORMAT JSON) + SELECT 1 FROM ONLY tool_calls x WHERE runtime_lock_id=$1 FOR KEY SHARE OF x +`).ReadAll() + require.NoError(t, err) + require.Len(t, results, 1) + assert.Contains(t, string(results[0].Rows[0][0]), "tool_calls_runtime_reference_idx") + t.Logf("history=%d RI probe: %s", size, raw) + assert.NotContains(t, string(raw), `"Node Type": "Seq Scan"`) + assert.LessOrEqual(t, planBuffers(t, raw), 10) + release, err := tx.Begin(t.Context()) + require.NoError(t, err) + require.NoError(t, release.QueryRow(t.Context(), ` +EXPLAIN (ANALYZE,BUFFERS,FORMAT JSON) DELETE FROM agent_runtime_locks WHERE agent_id=$1 AND id=$2 +`, agent, lease).Scan(&raw)) + t.Logf("history=%d delete: %s", size, raw) + assert.LessOrEqual(t, planBuffers(t, raw), 20) + require.NoError(t, release.Rollback(t.Context())) + } + _, err = tx.Exec(t.Context(), ` +WITH call AS (INSERT INTO tool_calls(agent_id,model_output_id,provider_call_id,name,input,type,state, +created_at) + VALUES($1,$2,'live','tool','{}','built_in','awaiting_authorization',now()) RETURNING id,created_at) +INSERT INTO content_blocks(agent_id,owner_kind,owner_model_output_id,ordinal,block_kind,tool_call_id, +created_at) + SELECT $1,'model_output',$2,100001,'tool_call',id,created_at FROM call +`, agent, output) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +UPDATE tool_calls SET state='ready' WHERE agent_id=$1 AND provider_call_id='live' +`, agent) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +UPDATE tool_calls SET state='running',runtime_lock_id=$2 WHERE agent_id=$1 AND provider_call_id='live' +`, agent, lease) + require.NoError(t, err) + var raw []byte + query := generatedQuery(t, "internal/executiondb/commands_runtime.sql.go", "ListExecutionUnfinishedTools") + require.NoError(t, tx.QueryRow(t.Context(), "EXPLAIN (ANALYZE,BUFFERS,FORMAT JSON) "+query, agent, + nil, lease).Scan(&raw)) + t.Logf("running tools plan=%s", raw) + require.LessOrEqual(t, planBuffers(t, raw), 20) + require.NotContains(t, string(raw), `"Node Type": "Seq Scan"`) + _, err = tx.Exec(t.Context(), `DELETE FROM agent_runtime_locks WHERE agent_id=$1 AND id=$2`, agent, lease) + require.ErrorContains(t, err, "tool_calls_runtime_lock_id_fkey") +} + +func planBuffers(t *testing.T, raw []byte) int { + t.Helper() + var plans []struct { + Plan struct { + Hit int `json:"Shared Hit Blocks"` + Read int `json:"Shared Read Blocks"` + } + } + require.NoError(t, json.Unmarshal(raw, &plans)) + return plans[0].Plan.Hit + plans[0].Plan.Read +} + +func TestExecutionHotPathScopeConstraints(t *testing.T) { + f, agent, lease := commandFixture(t) + _, err := f.pool.Exec(t.Context(), `UPDATE orgs SET idempotency_key=NULL; UPDATE projects SET idempotency_key=NULL`) + require.NoError(t, err) + other := seedExecutionFixture(t, f.pool) + project, config := uuid.New(), uuid.New() + _, err = f.pool.Exec(t.Context(), ` +INSERT INTO projects(id,org_id,name,created_at,updated_at) VALUES($1,$2,'other',now(),now()) +`, project, f.org) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), ` +INSERT INTO agent_configs SELECT (jsonb_populate_record(NULL::agent_configs, +to_jsonb(c)||jsonb_build_object('id',$2::uuid,'project_id',$3::uuid))).* FROM agent_configs c WHERE +id=$1 +`, f.config, config, project) + require.NoError(t, err) + u, h := f.handle(t, agent) + tool := toolCommand(t, h, lease, "built_in") + process := processOwner(t, f, u, agent, lease, tool.ID, "queued") + action := uuid.New() + _, err = u.DB().Exec(t.Context(), ` +INSERT INTO process_actions(id,org_id,project_id,agent_id,process_id,tool_call_id,runtime_lock_id, +action_kind,seq,created_at,updated_at) + VALUES($1,$2,$3,$4,$5,$6,$7,'read',1,now(),now()) +`, action, f.org, f.project, agent, process, tool.ID, lease) + require.NoError(t, err) + require.NoError(t, u.Commit(t.Context(), "scope fixture")) + for _, scope := range []struct { + name string + org, project, config, revision uuid.UUID + }{ + {"same", f.org, f.project, f.config, f.revision}, + {"other_org", other.org, other.project, other.config, other.revision}, + {"wrong_org", other.org, f.project, f.config, other.revision}, + {"other_project", f.org, project, config, f.revision}, + } { + t.Run("model/"+scope.name, func(t *testing.T) { + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + _, err = tx.Exec(t.Context(), ` +INSERT INTO model_call_contexts(id,org_id,project_id,agent_id,turn_id,operation_kind,attempt_number, +agent_config_id,configured_model_revision_id,input_event_sequence,runtime_lock_id,created_at, +opening_input_ids,opening_event_sequence) +SELECT $1,$2,$3,c.agent_id,c.turn_id,'normal',2,$4,$5,c.input_event_sequence,c.runtime_lock_id,now(), +c.opening_input_ids,c.opening_event_sequence + FROM model_call_contexts c WHERE agent_id=$6 +`, uuid.New(), scope.org, scope.project, scope.config, scope.revision, agent) + if scope.name == "same" { + require.NoError(t, err) + } else { + var constraint *pgconn.PgError + require.ErrorAs(t, err, &constraint) + require.Equal(t, "23503", constraint.Code) + require.Equal(t, "model_call_contexts_agent_scope_fkey", constraint.ConstraintName) + } + }) + } + for _, tc := range []struct { + table string + id uuid.UUID + }{{"processes", process}, {"process_actions", action}} { + for _, scope := range []struct { + name string + org, project uuid.UUID + }{{"org", other.org, f.project}, {"project", f.org, project}} { + t.Run(tc.table+"/"+scope.name, func(t *testing.T) { + _, err := f.pool.Exec(t.Context(), "UPDATE "+tc.table+ + " SET org_id=$1,project_id=$2 WHERE id=$3", scope.org, scope.project, tc.id) + var constraint *pgconn.PgError + require.ErrorAs(t, err, &constraint) + require.Equal(t, "23503", constraint.Code) + if scope.name == "org" { + require.Equal(t, tc.table+"_agent_scope_fkey", constraint.ConstraintName) + } + }) + } + } + for _, scope := range []struct { + name string + project uuid.UUID + }{{"org", other.project}, {"project", project}} { + t.Run("input/"+scope.name, func(t *testing.T) { + _, err := f.pool.Exec(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,queued_at) + VALUES($1,$2,'received','content','queued',now()) +`, scope.project, agent) + var constraint *pgconn.PgError + require.ErrorAs(t, err, &constraint) + require.Equal(t, "23503", constraint.Code) + require.Equal(t, "agent_inputs_project_id_agent_id_fkey", constraint.ConstraintName) + }) + } + var n int + require.NoError(t, f.pool.QueryRow(t.Context(), ` +SELECT count(*) FROM pg_constraint WHERE conname=ANY($1::text[]) AND convalidated +`, + []string{"model_call_contexts_agent_scope_fkey", "processes_agent_scope_fkey", + "process_actions_agent_scope_fkey"}).Scan(&n)) + require.Equal(t, 3, n) + require.NoError(t, f.pool.QueryRow(t.Context(), `SELECT count(*) FROM pg_constraint WHERE conname=ANY($1::text[])`, + []string{"model_call_contexts_org_id_fkey", "model_call_contexts_org_id_project_id_fkey", + "model_call_contexts_project_id_agent_id_fkey", "agent_inputs_project_id_fkey", + "processes_org_id_fkey", "processes_org_id_project_id_fkey", "processes_project_id_agent_id_fkey", + "process_actions_org_id_fkey", "process_actions_org_id_project_id_fkey"}).Scan(&n)) + require.Zero(t, n) +} diff --git a/migrations/000052_agent_execution_state.sql b/migrations/000052_agent_execution_state.sql index 933dc1180..3a6f1a7e6 100644 --- a/migrations/000052_agent_execution_state.sql +++ b/migrations/000052_agent_execution_state.sql @@ -449,3 +449,69 @@ DROP FUNCTION agent_turn_opening_content_inputs(uuid,uuid,bigint); DROP FUNCTION agent_turn_id_at_event_sequence(uuid,bigint); ALTER TABLE agent_wakeups DROP COLUMN metadata; + +CREATE INDEX tool_calls_runtime_reference_idx + ON tool_calls(runtime_lock_id) WHERE runtime_lock_id IS NOT NULL; + +DROP INDEX tool_calls_runtime_lock_idx; + +CREATE INDEX agent_inputs_received_content_idx + ON agent_inputs(agent_id, input_rank, queued_at, id) + WHERE state = 'received' AND input_kind = 'content'; + +CREATE UNIQUE INDEX agents_org_project_id_key + ON agents(org_id, project_id, id); + +ALTER TABLE model_call_contexts ADD CONSTRAINT model_call_contexts_agent_scope_fkey + FOREIGN KEY (org_id, project_id, agent_id) REFERENCES agents(org_id, project_id, id); + +ALTER TABLE processes ADD CONSTRAINT processes_agent_scope_fkey + FOREIGN KEY (org_id, project_id, agent_id) REFERENCES agents(org_id, project_id, id); + +ALTER TABLE process_actions ADD CONSTRAINT process_actions_agent_scope_fkey + FOREIGN KEY (org_id, project_id, agent_id) REFERENCES agents(org_id, project_id, id); + +ALTER TABLE model_call_contexts + DROP CONSTRAINT model_call_contexts_org_id_fkey, + DROP CONSTRAINT model_call_contexts_org_id_project_id_fkey, + DROP CONSTRAINT model_call_contexts_project_id_agent_id_fkey; + +ALTER TABLE agent_inputs DROP CONSTRAINT agent_inputs_project_id_fkey; + +ALTER TABLE processes + DROP CONSTRAINT processes_org_id_fkey, + DROP CONSTRAINT processes_org_id_project_id_fkey, + DROP CONSTRAINT processes_project_id_agent_id_fkey; + +ALTER TABLE process_actions + DROP CONSTRAINT process_actions_org_id_fkey, + DROP CONSTRAINT process_actions_org_id_project_id_fkey; + +-- +goose StatementBegin +CREATE FUNCTION agent_payload_created_at( + p_project_id uuid, p_agent_id uuid, p_kind text, p_id uuid, p_admission boolean +) +RETURNS timestamptz +LANGUAGE plpgsql STABLE +AS $$ +DECLARE result timestamptz; +BEGIN + IF NOT EXISTS (SELECT 1 FROM agents WHERE project_id = p_project_id AND id = p_agent_id) THEN + RETURN NULL; + END IF; + CASE p_kind + WHEN 'agent_input' THEN + SELECT CASE WHEN p_admission THEN statement_timestamp() ELSE queued_at END INTO result + FROM agent_inputs WHERE agent_id = p_agent_id AND id = p_id; + WHEN 'model_output' THEN + SELECT created_at INTO result FROM model_outputs WHERE agent_id = p_agent_id AND id = p_id; + WHEN 'tool_call_result' THEN + SELECT completed_at INTO result FROM tool_call_results WHERE agent_id = p_agent_id AND id = p_id; + WHEN 'context_checkpoint' THEN + SELECT created_at INTO result FROM context_checkpoints WHERE agent_id = p_agent_id AND id = p_id; + ELSE RETURN NULL; + END CASE; + RETURN result; +END; +$$; +-- +goose StatementEnd diff --git a/migrations/000053_agent_payload_query_plans.sql b/migrations/000053_agent_payload_query_plans.sql deleted file mode 100644 index e1ce24670..000000000 --- a/migrations/000053_agent_payload_query_plans.sql +++ /dev/null @@ -1,30 +0,0 @@ --- +goose Up - --- +goose StatementBegin -CREATE FUNCTION agent_payload_created_at( - p_project_id uuid, p_agent_id uuid, p_kind text, p_id uuid, p_admission boolean -) -RETURNS timestamptz -LANGUAGE plpgsql STABLE -AS $$ -DECLARE result timestamptz; -BEGIN - IF NOT EXISTS (SELECT 1 FROM agents WHERE project_id = p_project_id AND id = p_agent_id) THEN - RETURN NULL; - END IF; - CASE p_kind - WHEN 'agent_input' THEN - SELECT CASE WHEN p_admission THEN statement_timestamp() ELSE queued_at END INTO result - FROM agent_inputs WHERE agent_id = p_agent_id AND id = p_id; - WHEN 'model_output' THEN - SELECT created_at INTO result FROM model_outputs WHERE agent_id = p_agent_id AND id = p_id; - WHEN 'tool_call_result' THEN - SELECT completed_at INTO result FROM tool_call_results WHERE agent_id = p_agent_id AND id = p_id; - WHEN 'context_checkpoint' THEN - SELECT created_at INTO result FROM context_checkpoints WHERE agent_id = p_agent_id AND id = p_id; - ELSE RETURN NULL; - END CASE; - RETURN result; -END; -$$; --- +goose StatementEnd From ba9ab0f67c829208f8201da69c977da09106e294 Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:55:10 -0700 Subject: [PATCH 08/19] perf: move, rank and admit backlog inputs with indexed lookups - Move a queued input by reading it and its anchor, finding the destination neighbor with one indexed lookup, and updating one row. When a rank gap is exhausted, renumber that agent's pending inputs in the mode with a single statement and retry once. - Rank new content inputs from the last pending input in their mode; config changes skip ranking. - Admit only the selected queued input or pending steering inputs, with the remaining-input flags in the same statement. - Serve every engine backlog read from one partial index over pending content inputs, replacing the two project-scoped per-mode indexes. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../execution_cutover_integration_test.go | 6 +- .../backlog_cost_integration_test.go | 190 ++++++++++++++ .../agentexecution/commands_backlog.go | 184 ++++++++++++++ .../commands_backlog_integration_test.go | 199 ++++++++++++++- .../agentexecution/commands_inputs.go | 188 +------------- .../hot_path_plans_integration_test.go | 10 +- .../internal/executiondb/batch.go | 10 +- .../executiondb/commands_inputs.sql.go | 232 ++++++++++++++++-- .../load_count_integration_test.go | 23 +- .../queries/commands_inputs.sql | 64 ++++- .../internal/agentexecution/queries/load.sql | 10 +- migrations/000052_agent_execution_state.sql | 5 +- 12 files changed, 892 insertions(+), 229 deletions(-) create mode 100644 internal/storage/internal/agentexecution/backlog_cost_integration_test.go create mode 100644 internal/storage/internal/agentexecution/commands_backlog.go diff --git a/internal/dbmigrate/execution_cutover_integration_test.go b/internal/dbmigrate/execution_cutover_integration_test.go index 3f4bc0415..bddfb6c88 100644 --- a/internal/dbmigrate/execution_cutover_integration_test.go +++ b/internal/dbmigrate/execution_cutover_integration_test.go @@ -131,7 +131,7 @@ SELECT count(*) FROM pg_constraint WHERE conname=ANY($1::text[]) unique bool }{ {"tool_calls_runtime_reference_idx", "USING btree (runtime_lock_id) WHERE (runtime_lock_id IS NOT NULL)", false}, - {"agent_inputs_received_content_idx", "USING btree (agent_id, input_rank, queued_at, id)", false}, + {"agent_inputs_received_content_idx", "USING btree (agent_id, delivery_mode, input_rank, queued_at, id)", false}, {"agents_org_project_id_key", "USING btree (org_id, project_id, id)", true}, } { var definition string @@ -151,7 +151,9 @@ SELECT pg_get_indexdef(indexrelid),indisready,indisvalid,indisunique FROM pg_ind var payloadInstalled, oldIndexAbsent bool require.NoError(t, pool.QueryRow(ctx, ` SELECT to_regprocedure('agent_payload_created_at(uuid,uuid,text,uuid,boolean)') IS NOT NULL, - to_regclass('tool_calls_runtime_lock_idx') IS NULL + to_regclass('tool_calls_runtime_lock_idx') IS NULL AND + to_regclass('agent_inputs_content_queued_idx') IS NULL AND + to_regclass('agent_inputs_content_steering_idx') IS NULL `).Scan(&payloadInstalled, &oldIndexAbsent)) require.True(t, payloadInstalled) require.True(t, oldIndexAbsent) diff --git a/internal/storage/internal/agentexecution/backlog_cost_integration_test.go b/internal/storage/internal/agentexecution/backlog_cost_integration_test.go new file mode 100644 index 000000000..4b28a2c12 --- /dev/null +++ b/internal/storage/internal/agentexecution/backlog_cost_integration_test.go @@ -0,0 +1,190 @@ +//go:build integration + +package agentexecution_test + +import ( + "encoding/json" + "strconv" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestExecutionBacklogCosts(t *testing.T) { + f, agent, _ := commandFixture(t) + other := f.create(t, uuid.New(), uuid.Nil) + previous := 0 + for _, size := range []int{100, 10000} { + _, err := f.pool.Exec(t.Context(), ` +INSERT INTO agent_inputs(project_id,agent_id,state,input_kind,delivery_mode,input_rank,queued_at,canceled_at) +SELECT $1,a.id,CASE WHEN history THEN 'canceled' ELSE 'received' END,'content',mode,n*1024,now(), + CASE WHEN history THEN now() END +FROM (VALUES($2::uuid),($3::uuid)) a(id) +CROSS JOIN (VALUES('queued'),('steering')) m(mode) +CROSS JOIN (VALUES(false),(true)) h(history) +CROSS JOIN LATERAL generate_series($4::int*CASE WHEN history THEN 10 ELSE 1 END+1, + $5::int*CASE WHEN history THEN 10 ELSE 1 END) n +WHERE mode='queued' OR a.id=$3 +`, f.project, agent, other, previous, size) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), ` +INSERT INTO content_blocks(agent_id,owner_kind,owner_agent_input_id,ordinal,block_kind,text_content,created_at) +SELECT i.agent_id,'agent_input',i.id,0,'text','pending',i.queued_at FROM agent_inputs i +WHERE i.agent_id=$1 AND i.state='received' AND i.input_rank>$2 +`, agent, previous*1024) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), `ANALYZE agent_inputs; ANALYZE agents; ANALYZE content_blocks`) + require.NoError(t, err) + previous = size + var first, second, last uuid.UUID + require.NoError(t, f.pool.QueryRow(t.Context(), ` +SELECT (array_agg(id ORDER BY input_rank))[1],(array_agg(id ORDER BY input_rank))[2], + (array_agg(id ORDER BY input_rank DESC))[1] FROM agent_inputs WHERE agent_id=$1 AND state='received' +`, agent).Scan(&first, &second, &last)) + if size == 10000 { + t.Run("index_consumers", func(t *testing.T) { + conn, err := f.pool.Acquire(t.Context()) + require.NoError(t, err) + defer conn.Release() + for _, consumer := range []struct{ path, name string }{ + {"../dbsqlc/agent_inputs.sql.go", "ListQueuedBacklogInputs"}, + {"internal/executiondb/commands_runtime.sql.go", "CancelExecutionInputs"}, + } { + query := generatedQuery(t, consumer.path, consumer.name) + plans, err := conn.Conn().PgConn().Exec(t.Context(), "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+query).ReadAll() + require.NoError(t, err) + require.Len(t, plans, 1) + generic := string(plans[0].Rows[0][0]) + t.Logf("generic %s %s", consumer.name, generic) + require.Contains(t, generic, "agent_inputs_received_content_idx") + require.NotContains(t, generic, "agent_inputs_content_queued_idx") + require.NotContains(t, generic, "agent_inputs_content_steering_idx") + } + }) + } + for _, command := range []string{"move", "enqueue", "admit", "renumber"} { + t.Run(command+"/"+strconv.Itoa(size), func(t *testing.T) { + if command == "renumber" { + _, err = f.pool.Exec(t.Context(), `UPDATE agent_inputs SET input_rank=1 WHERE agent_id=$1 AND id=$2`, agent, first) + require.NoError(t, err) + _, err = f.pool.Exec(t.Context(), `UPDATE agent_inputs SET input_rank=2 WHERE agent_id=$1 AND id=$2`, + agent, second) + require.NoError(t, err) + defer func() { + _, err := f.pool.Exec(t.Context(), `UPDATE agent_inputs SET input_rank=CASE id WHEN $2 THEN 1024 ELSE 2048 END +WHERE agent_id=$1 AND id IN ($2,$3)`, agent, first, second) + require.NoError(t, err) + }() + } + u, h := f.handle(t, agent) + defer func() { _ = u.Rollback(t.Context()) }() + started := time.Now() + trace, err := agentexecution.TraceExecutionCommand(u, func() error { + switch command { + case "enqueue": + _, err := h.ReceiveContent(t.Context(), agentexecution.ReceiveContentInput{ + DeliveryMode: "queued", Content: []agentexecution.Content{{Kind: "text", Text: "new"}}, + }) + return err + case "admit": + admitted, err := h.AdmitInputs(t.Context()) + if err == nil { + require.Len(t, admitted.Inputs, 1) + require.Equal(t, first, admitted.Inputs[0].ID) + } + return err + default: + _, err := h.ChangeBacklog(t.Context(), agentexecution.BacklogChange{ + ID: last, DeliveryMode: "queued", Position: "before", BeforeID: second}) + return err + } + }) + require.NoError(t, err) + elapsed := time.Since(started) + bound := 4 + if command == "admit" || command == "renumber" { + bound = 7 + } + assert.LessOrEqual(t, len(trace), bound) + if command == "renumber" { + rows, err := u.DB().Query(t.Context(), `SELECT agent_id,delivery_mode,count(*) FROM agent_inputs +WHERE xmin=pg_current_xact_id()::text::xid GROUP BY agent_id,delivery_mode`) + require.NoError(t, err) + var groups int + for rows.Next() { + var changedAgent uuid.UUID + var mode string + var count int + require.NoError(t, rows.Scan(&changedAgent, &mode, &count)) + require.Equal(t, agent, changedAgent) + require.Equal(t, "queued", mode) + require.Equal(t, size, count) + groups++ + } + require.NoError(t, rows.Err()) + rows.Close() + require.Equal(t, 1, groups) + } + require.NoError(t, u.Rollback(t.Context())) + t.Logf("pending=%d history=%d command=%s statements=%d duration=%s", size, size*10, command, len(trace), elapsed) + if command == "renumber" { + var bulk int + for _, statement := range trace { + if statement.Affected > 1 { + bulk++ + assert.EqualValues(t, size, statement.Affected) + } + } + require.Equal(t, 1, bulk) + require.LessOrEqual(t, len(trace), 7) + } + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + _, err = tx.Exec(t.Context(), `SET LOCAL plan_cache_mode=force_generic_plan`) + require.NoError(t, err) + var total int + for _, statement := range trace { + var raw []byte + require.NoError(t, tx.QueryRow(t.Context(), + "EXPLAIN (ANALYZE,BUFFERS,FORMAT JSON) "+statement.Query.SQL, statement.Query.Arguments...).Scan(&raw)) + name, _, _ := strings.Cut(statement.Query.SQL, "\n") + t.Logf("%s %s", name, raw) + if strings.Contains(name, "BacklogRank") || strings.Contains(name, "SelectExecutionAdmission") || + strings.Contains(name, "ReceiveExecutionInput") { + plans, err := tx.Conn().PgConn().Exec(t.Context(), + "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+statement.Query.SQL).ReadAll() + require.NoError(t, err) + require.Len(t, plans, 1) + generic := string(plans[0].Rows[0][0]) + require.Contains(t, generic, "agent_inputs_received_content_idx") + if !strings.Contains(name, "ReceiveExecutionInput") { + require.NotContains(t, generic, `"Node Type": "Seq Scan"`) + } + t.Logf("generic %s %s", name, generic) + } + total += planBuffers(t, raw) + if command != "renumber" { + assert.LessOrEqual(t, planBuffers(t, raw), 128, name) + var plan []struct { + Plan struct { + Rows float64 `json:"Actual Rows"` + } + } + require.NoError(t, json.Unmarshal(raw, &plan)) + assert.LessOrEqual(t, plan[0].Plan.Rows, float64(3), name) + } + } + t.Logf("command=%s total_buffers=%d", command, total) + if command != "renumber" { + assert.LessOrEqual(t, total, 256) + } + }) + } + } +} diff --git a/internal/storage/internal/agentexecution/commands_backlog.go b/internal/storage/internal/agentexecution/commands_backlog.go new file mode 100644 index 000000000..123dec4e1 --- /dev/null +++ b/internal/storage/internal/agentexecution/commands_backlog.go @@ -0,0 +1,184 @@ +package agentexecution + +import ( + "context" + "errors" + "math" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution/internal/executiondb" + "github.com/omnara-ai/omnara/internal/storage/storeerr" +) + +type BacklogChange struct { + ID uuid.UUID + DeliveryMode string + Cancel bool + BeforeID uuid.UUID + Position string +} + +func (h *Handle) ChangeBacklog(ctx context.Context, change BacklogChange) (bool, error) { + return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { + if change.ID == uuid.Nil || change.DeliveryMode != "queued" && change.DeliveryMode != "steering" { + return false, errors.New("invalid backlog change") + } + switch change.Position { + case "", "front", "back": + if change.BeforeID != uuid.Nil { + return false, errors.New("backlog change does not accept an anchor") + } + case "before", "after": + if change.BeforeID == uuid.Nil || change.BeforeID == change.ID { + return false, storeerr.ErrStateTransitionConflict + } + default: + return false, errors.New("invalid backlog position") + } + if change.Cancel && change.Position != "" || + (change.Cancel || change.Position != "") && change.DeliveryMode != "queued" { + return false, errors.New("only queued backlog supports cancellation and movement") + } + + selected, anchor, err := h.backlogInputs(ctx, change) + if err != nil { + return false, err + } + if selected.InputKind != "content" || selected.State != "received" { + if change.DeliveryMode == "steering" && selected.InputKind == "content" && + selected.State == "resolved" && selected.AdmittedEventID != nil { + return false, nil + } + return false, storeerr.ErrStateTransitionConflict + } + if change.Cancel { + if selected.DeliveryMode != "queued" { + return false, storeerr.ErrStateTransitionConflict + } + return h.changeBacklogRow(ctx, m, selected, true) + } + if change.Position != "" && selected.DeliveryMode != "queued" { + return false, storeerr.ErrStateTransitionConflict + } + if change.Position == "" && change.DeliveryMode == selected.DeliveryMode { + if change.DeliveryMode == "steering" { + return false, nil + } + return false, storeerr.ErrStateTransitionConflict + } + var changed bool + for attempt := range 2 { + rank, available, err := h.backlogDestination(ctx, change, anchor) + if err != nil { + return false, err + } + if available { + selected.DeliveryMode = change.DeliveryMode + selected.InputRank = rank + updated, err := h.changeBacklogRow(ctx, m, selected, false) + return changed || updated, err + } + if attempt == 1 { + break + } + n, err := executiondb.New().RebalanceExecutionBacklog(ctx, h.unit.DB(), executiondb.RebalanceExecutionBacklogParams{ + AgentID: h.route.AgentID, Mode: change.DeliveryMode, + }) + if err != nil { + return false, err + } + if n > 0 { + changed = true + m.changed() + } + selected, anchor, err = h.backlogInputs(ctx, change) + if err != nil { + return false, err + } + } + return false, storeerr.ErrStateTransitionConflict + }) +} + +func (h *Handle) backlogInputs( + ctx context.Context, + change BacklogChange, +) (executiondb.GetExecutionBacklogInputsRow, executiondb.GetExecutionBacklogInputsRow, error) { + var selected, anchor executiondb.GetExecutionBacklogInputsRow + ids := []uuid.UUID{change.ID} + if change.BeforeID != uuid.Nil { + ids = append(ids, change.BeforeID) + } + rows, err := executiondb.New().GetExecutionBacklogInputs(ctx, h.unit.DB(), + executiondb.GetExecutionBacklogInputsParams{AgentID: h.route.AgentID, Ids: ids}) + if err != nil { + return selected, anchor, err + } + for _, row := range rows { + if row.ID == change.ID { + selected = row + } + if row.ID == change.BeforeID { + anchor = row + } + } + if selected.ID == uuid.Nil || change.BeforeID != uuid.Nil && + (anchor.ID == uuid.Nil || anchor.State != "received" || + anchor.InputKind != "content" || anchor.DeliveryMode != "queued") { + return selected, anchor, storeerr.ErrStateTransitionConflict + } + return selected, anchor, nil +} + +func (h *Handle) backlogDestination( + ctx context.Context, + change BacklogChange, + anchor executiondb.GetExecutionBacklogInputsRow, +) (int64, bool, error) { + q := executiondb.New() + var lower, upper int64 + var err error + switch change.Position { + case "front": + upper, err = q.FirstExecutionBacklogRank(ctx, h.unit.DB(), executiondb.FirstExecutionBacklogRankParams{ + AgentID: h.route.AgentID, Mode: change.DeliveryMode, SkipID: change.ID, + }) + case "before": + upper = anchor.InputRank + lower, err = q.PreviousExecutionBacklogRank(ctx, h.unit.DB(), executiondb.PreviousExecutionBacklogRankParams{ + AgentID: h.route.AgentID, Mode: change.DeliveryMode, SkipID: change.ID, + Rank: anchor.InputRank, QueuedAt: anchor.QueuedAt, AnchorID: anchor.ID, + }) + case "after": + lower = anchor.InputRank + upper, err = q.NextExecutionBacklogRank(ctx, h.unit.DB(), executiondb.NextExecutionBacklogRankParams{ + AgentID: h.route.AgentID, Mode: change.DeliveryMode, SkipID: change.ID, + Rank: anchor.InputRank, QueuedAt: anchor.QueuedAt, AnchorID: anchor.ID, + }) + default: + lower, err = q.LastExecutionBacklogRank(ctx, h.unit.DB(), executiondb.LastExecutionBacklogRankParams{ + AgentID: h.route.AgentID, Mode: change.DeliveryMode, SkipID: change.ID, + }) + } + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return 0, false, err + } + if upper > 0 { + return lower + (upper-lower)/2, upper-lower > 1, nil + } + if lower > math.MaxInt64-1024 { + return 0, false, errors.New("backlog rank overflow") + } + return lower + 1024, true, nil +} + +func (h *Handle) changeBacklogRow(ctx context.Context, m *executionMutation, + row executiondb.GetExecutionBacklogInputsRow, cancel bool) (bool, error) { + n, err := executiondb.New().ChangeExecutionInput(ctx, h.unit.DB(), executiondb.ChangeExecutionInputParams{ + AgentID: h.route.AgentID, ID: row.ID, Mode: row.DeliveryMode, Rank: row.InputRank, Cancel: cancel}) + if n > 0 { + m.changed() + } + return n > 0, err +} diff --git a/internal/storage/internal/agentexecution/commands_backlog_integration_test.go b/internal/storage/internal/agentexecution/commands_backlog_integration_test.go index f333c689d..866395240 100644 --- a/internal/storage/internal/agentexecution/commands_backlog_integration_test.go +++ b/internal/storage/internal/agentexecution/commands_backlog_integration_test.go @@ -3,12 +3,17 @@ package agentexecution_test import ( + "fmt" + "slices" "testing" + "time" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/omnara-ai/omnara/internal/storage/internal/agentexecution" + "github.com/omnara-ai/omnara/internal/storage/internal/lifecyclelock" "github.com/omnara-ai/omnara/internal/storage/storeerr" + "github.com/omnara-ai/omnara/internal/testutil/integrationdb" "github.com/stretchr/testify/require" ) @@ -89,7 +94,8 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { first := receiveCommand(t, h, "queued", "first") second := receiveCommand(t, h, "queued", "second") third := receiveCommand(t, h, "steering", "third") - promote := agentexecution.BacklogChange{ID: second.ID, DeliveryMode: "steering"} + promote := agentexecution.BacklogChange{ + ID: second.ID, DeliveryMode: "steering"} changed, err := h.ChangeBacklog(t.Context(), promote) require.NoError(t, err) require.True(t, changed) @@ -108,7 +114,8 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { assertCommandState(t, u, h) changed, err = h.ChangeBacklog( t.Context(), - agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "queued"}, + agentexecution.BacklogChange{ + ID: third.ID, DeliveryMode: "queued"}, ) require.NoError(t, err) require.True(t, changed) @@ -117,7 +124,8 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { u, h = f.handle(t, a) _, err = h.ChangeBacklog( t.Context(), - agentexecution.BacklogChange{ID: third.ID, DeliveryMode: "queued"}, + agentexecution.BacklogChange{ + ID: third.ID, DeliveryMode: "queued"}, ) require.NoError(t, u.Rollback(t.Context())) @@ -129,7 +137,8 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { require.NoError(t, err) require.False(t, changed) assertCommandState(t, u, h) - cancel := agentexecution.BacklogChange{ID: first.ID, DeliveryMode: "queued", Cancel: true} + cancel := agentexecution.BacklogChange{ + ID: first.ID, DeliveryMode: "queued", Cancel: true} changed, err = h.ChangeBacklog(t.Context(), cancel) require.NoError(t, err) require.True(t, changed) @@ -142,3 +151,185 @@ func TestExecutionBacklogDeliveryReplay(t *testing.T) { require.ErrorIs(t, err, storeerr.ErrStateTransitionConflict) assertCommandState(t, u, h) } + +func TestExecutionBacklogRepeatedGapRebalance(t *testing.T) { + f, a, _ := commandFixture(t) + other := f.create(t, uuid.New(), uuid.Nil) + ou, oh := f.handle(t, other) + receiveCommand(t, oh, "queued", "other queued") + receiveCommand(t, oh, "steering", "other steering") + otherBefore := readCommandBacklog(t, ou, other) + require.NoError(t, ou.Commit(t.Context(), "other backlog")) + u, h := f.handle(t, a) + ids := make([]uuid.UUID, 4) + for i := range ids { + ids[i] = receiveCommand(t, h, "queued", "queued").ID + } + steering := receiveCommand(t, h, "steering", "steering") + canceled := receiveCommand(t, h, "queued", "canceled") + _, err := h.ChangeBacklog(t.Context(), agentexecution.BacklogChange{ + ID: canceled.ID, DeliveryMode: "queued", Cancel: true}) + require.NoError(t, err) + var untouched []backlogRecord + for _, row := range readCommandBacklog(t, u, a) { + if row.ID == steering.ID || row.ID == canceled.ID { + untouched = append(untouched, row) + } + } + expected := append([]uuid.UUID(nil), ids...) + bulk := 0 + for iteration := range 11 { + selected := ids[2+iteration%2] + trace, err := agentexecution.TraceExecutionCommand(u, func() error { + changed, err := h.ChangeBacklog(t.Context(), agentexecution.BacklogChange{ + ID: selected, DeliveryMode: "queued", Position: "before", BeforeID: ids[1]}) + if err == nil { + require.True(t, changed) + } + return err + }) + require.NoError(t, err) + for _, statement := range trace { + if statement.Affected > 1 { + bulk++ + require.EqualValues(t, 4, statement.Affected) + } + } + expected = slices.DeleteFunc(expected, func(id uuid.UUID) bool { return id == selected }) + index := slices.Index(expected, ids[1]) + expected = slices.Insert(expected, index, selected) + var actual []uuid.UUID + var rank int64 + var remaining []backlogRecord + for _, row := range readCommandBacklog(t, u, a) { + if row.ID == steering.ID || row.ID == canceled.ID { + remaining = append(remaining, row) + continue + } + actual = append(actual, row.ID) + require.Greater(t, row.Rank, rank) + rank = row.Rank + } + require.Equal(t, expected, actual) + require.Equal(t, untouched, remaining) + } + require.Equal(t, 1, bulk) + require.Equal(t, otherBefore, readCommandBacklog(t, u, other)) + assertCommandState(t, u, h) + require.NoError(t, u.Commit(t.Context(), "rebalanced backlog")) +} + +func TestExecutionBacklogAnchorConcurrentInvalidation(t *testing.T) { + for _, operation := range []string{"cancel", "admit", "promote"} { + t.Run(operation, func(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + anchor := receiveCommand(t, h, "queued", "anchor") + moved := receiveCommand(t, h, "queued", "moved") + require.NoError(t, u.Commit(t.Context(), "backlog")) + u, h = f.handle(t, a) + waiting, err := agentexecution.BeginTestUnit(t.Context(), agentexecution.NewCell(f.pool, nil, nil)) + require.NoError(t, err) + defer func() { _ = waiting.Rollback(t.Context()) }() + plan, err := waiting.PlanAgents(t.Context(), []lifecyclelock.AgentRef{{ProjectID: f.project, AgentID: a}}, nil) + require.NoError(t, err) + result := make(chan error, 1) + go func() { + err := waiting.LockAgents(t.Context(), plan) + if err == nil { + var handle *agentexecution.Handle + handle, err = waiting.Agent(agentexecution.AgentRoute{ProjectID: f.project, AgentID: a, RootAgentID: a}) + if err == nil { + _, err = handle.ChangeBacklog(t.Context(), agentexecution.BacklogChange{ + ID: moved.ID, DeliveryMode: "queued", Position: "before", BeforeID: anchor.ID}) + } + } + result <- err + }() + integrationdb.WaitForLockWaiters(t, t.Context(), f.pool, "FROM agents", 1) + switch operation { + case "admit": + admitted, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + require.Len(t, admitted.Inputs, 1) + require.Equal(t, anchor.ID, admitted.Inputs[0].ID) + case "cancel": + _, err = h.ChangeBacklog(t.Context(), agentexecution.BacklogChange{ + ID: anchor.ID, DeliveryMode: "queued", Cancel: true}) + require.NoError(t, err) + case "promote": + _, err = h.ChangeBacklog(t.Context(), agentexecution.BacklogChange{ + ID: anchor.ID, DeliveryMode: "steering"}) + require.NoError(t, err) + } + before := readCommandBacklog(t, u, a) + require.NoError(t, u.Commit(t.Context(), "invalidate anchor")) + select { + case err := <-result: + require.ErrorIs(t, err, storeerr.ErrStateTransitionConflict) + case <-time.After(5 * time.Second): + t.Fatal("move did not finish") + } + require.NoError(t, waiting.Rollback(t.Context())) + check, _ := f.handle(t, a) + require.Equal(t, before, readCommandBacklog(t, check, a)) + }) + } +} + +func TestExecutionAdmissionRemainingInputs(t *testing.T) { + f, a, _ := commandFixture(t) + for _, queued := range []int{0, 1, 2} { + for _, steering := range []int{0, 1, 2} { + t.Run(fmt.Sprintf("queued=%d/steering=%d", queued, steering), func(t *testing.T) { + u, h := f.handle(t, a) + for range queued { + receiveCommand(t, h, "queued", "queued") + } + for range steering { + receiveCommand(t, h, "steering", "steering") + } + before, err := h.LoadExecution(t.Context()) + require.NoError(t, err) + admitted, err := h.AdmitInputs(t.Context()) + require.NoError(t, err) + expected := steering + if expected == 0 && queued > 0 { + expected = 1 + } + require.Len(t, admitted.Inputs, expected) + after := assertCommandState(t, u, h) + if expected == 0 { + require.Equal(t, before.View, after.View) + } + require.False(t, after.View.Inputs.Steering) + require.Equal(t, queued > 1 || steering > 0 && queued > 0, after.View.Inputs.Queued) + require.NoError(t, u.Rollback(t.Context())) + }) + } + } +} + +func TestExecutionConfigInputDoesNotUseBacklogRank(t *testing.T) { + f, a, _ := commandFixture(t) + u, h := f.handle(t, a) + for _, mode := range []string{"queued", "steering"} { + receiveCommand(t, h, mode, "content") + } + for range 2 { + config, err := h.ActivateConfig(t.Context(), agentexecution.ActivateConfigInput{ConfigID: f.config}) + require.NoError(t, err) + var rank int64 + require.NoError(t, u.DB().QueryRow(t.Context(), + "SELECT input_rank FROM agent_inputs WHERE agent_id=$1 AND id=$2", a, config.InputID).Scan(&rank)) + require.EqualValues(t, 1024, rank) + assertCommandState(t, u, h) + } + for _, mode := range []string{"queued", "steering"} { + next := receiveCommand(t, h, mode, "content") + var rank int64 + require.NoError(t, u.DB().QueryRow(t.Context(), + "SELECT input_rank FROM agent_inputs WHERE agent_id=$1 AND id=$2", a, next.ID).Scan(&rank)) + require.EqualValues(t, 2048, rank) + } +} diff --git a/internal/storage/internal/agentexecution/commands_inputs.go b/internal/storage/internal/agentexecution/commands_inputs.go index 076206476..40355849d 100644 --- a/internal/storage/internal/agentexecution/commands_inputs.go +++ b/internal/storage/internal/agentexecution/commands_inputs.go @@ -4,8 +4,6 @@ import ( "context" "encoding/json" "errors" - "math" - "slices" "github.com/google/uuid" "github.com/jackc/pgx/v5" @@ -135,20 +133,20 @@ func (h *Handle) admitInputs( kind AdmissionKind, ) (InputAdmission, error) { q := executiondb.New() - rows, err := q.ListExecutionBacklog( - ctx, - h.unit.DB(), - executiondb.ListExecutionBacklogParams{AgentID: h.route.AgentID}, - ) + mode := "steering" + var limit *int64 + if kind == AdmitOneQueued { + mode = "queued" + limit = new(int64(1)) + } + rows, err := q.SelectExecutionAdmission(ctx, h.unit.DB(), executiondb.SelectExecutionAdmissionParams{ + AgentID: h.route.AgentID, Mode: mode, RowLimit: limit, + }) if err != nil { return InputAdmission{}, err } result := InputAdmission{TurnID: uuid.New()} for _, row := range rows { - if kind == AdmitAllSteering && row.DeliveryMode != "steering" || - kind == AdmitOneQueued && row.DeliveryMode != "queued" { - continue - } event, err := h.appendEvent(ctx, result.TurnID, "agent_input", row.ID, uuid.Nil, uuid.Nil, true) if err != nil { return InputAdmission{}, err @@ -170,9 +168,6 @@ func (h *Handle) admitInputs( Event: event, }, ) - if kind == AdmitOneQueued { - break - } } if len(result.Inputs) == 0 { return InputAdmission{}, nil @@ -205,19 +200,8 @@ func (h *Handle) admitInputs( turn.InitialOpening.InputIDs = append(turn.InitialOpening.InputIDs, row.ID) } loaded.View = ExecutionView{AgentID: h.route.AgentID, State: loaded.View.State, Turn: turn} - admitted := make(map[uuid.UUID]bool, len(result.Inputs)) - for _, input := range result.Inputs { - admitted[input.ID] = true - } - for _, row := range rows { - if admitted[row.ID] { - continue - } - if row.DeliveryMode == "steering" { - loaded.View.Inputs.Steering = true - } else { - loaded.View.Inputs.Queued = true - } + loaded.View.Inputs = InputAvailability{ + Steering: valueOrZero(rows[0].RemainingSteering), Queued: rows[0].RemainingQueued, } loaded.databaseNow = last.Time if err := m.updated(*loaded); err != nil { @@ -225,153 +209,3 @@ func (h *Handle) admitInputs( } return result, nil } - -type BacklogChange struct { - ID uuid.UUID - DeliveryMode string - Cancel bool - BeforeID uuid.UUID - Position string -} - -func (h *Handle) ChangeBacklog(ctx context.Context, change BacklogChange) (bool, error) { - return executeCommand(ctx, h, func(m *executionMutation) (bool, error) { - if change.ID == uuid.Nil || change.DeliveryMode != "queued" && change.DeliveryMode != "steering" { - return false, errors.New("invalid backlog change") - } - switch change.Position { - case "", "front", "back": - if change.BeforeID != uuid.Nil { - return false, errors.New("backlog change does not accept an anchor") - } - case "before", "after": - if change.BeforeID == uuid.Nil || change.BeforeID == change.ID { - return false, storeerr.ErrStateTransitionConflict - } - default: - return false, errors.New("invalid backlog position") - } - if change.Cancel && change.Position != "" || - (change.Cancel || change.Position != "") && change.DeliveryMode != "queued" { - return false, errors.New("only queued backlog supports cancellation and movement") - } - q := executiondb.New() - rows, err := q.ListExecutionBacklog(ctx, h.unit.DB(), - executiondb.ListExecutionBacklogParams{AgentID: h.route.AgentID}) - if err != nil { - return false, err - } - index := slices.IndexFunc( - rows, - func(row executiondb.ListExecutionBacklogRow) bool { return row.ID == change.ID }, - ) - if index < 0 { - if change.DeliveryMode == "steering" { - row, err := q.FindExecutionInput(ctx, h.unit.DB(), - executiondb.FindExecutionInputParams{AgentID: h.route.AgentID, ID: change.ID}) - if err == nil && row.InputKind == "content" && row.State == "resolved" && - row.AdmittedEventID != nil { - return false, nil - } - if err != nil && !errors.Is(err, pgx.ErrNoRows) { - return false, err - } - } - return false, storeerr.ErrStateTransitionConflict - } - selected := rows[index] - if change.Cancel { - if selected.DeliveryMode != "queued" { - return false, storeerr.ErrStateTransitionConflict - } - return h.changeBacklogRow(ctx, m, selected, true) - } - if change.Position != "" && selected.DeliveryMode != "queued" { - return false, storeerr.ErrStateTransitionConflict - } - if change.Position == "" && change.DeliveryMode == selected.DeliveryMode { - if change.DeliveryMode == "steering" { - return false, nil - } - return false, storeerr.ErrStateTransitionConflict - } - rows = slices.DeleteFunc(rows, func(row executiondb.ListExecutionBacklogRow) bool { - return row.DeliveryMode != change.DeliveryMode - }) - rank, available, err := backlogRank(rows, change) - if err != nil { - return false, err - } - var changed bool - if !available { - for i := range rows { - rows[i].InputRank = int64(i+1) * 1024 - updated, err := h.changeBacklogRow(ctx, m, rows[i], false) - if err != nil { - return false, err - } - changed = changed || updated - } - rank, available, err = backlogRank(rows, change) - if err != nil { - return false, err - } - if !available { - return false, storeerr.ErrStateTransitionConflict - } - } - selected.DeliveryMode = change.DeliveryMode - selected.InputRank = rank - updated, err := h.changeBacklogRow(ctx, m, selected, false) - return changed || updated, err - }) -} - -func (h *Handle) changeBacklogRow(ctx context.Context, m *executionMutation, - row executiondb.ListExecutionBacklogRow, cancel bool) (bool, error) { - n, err := executiondb.New().ChangeExecutionInput(ctx, h.unit.DB(), executiondb.ChangeExecutionInputParams{ - AgentID: h.route.AgentID, ID: row.ID, Mode: row.DeliveryMode, Rank: row.InputRank, Cancel: cancel}) - if n > 0 { - m.changed() - } - return n > 0, err -} - -func backlogRank(rows []executiondb.ListExecutionBacklogRow, change BacklogChange) (int64, bool, error) { - neighbors := make([]int64, 0, len(rows)) - anchor := -1 - for _, row := range rows { - if row.ID == change.ID { - continue - } - if row.ID == change.BeforeID { - anchor = len(neighbors) - } - neighbors = append(neighbors, row.InputRank) - } - destination := len(neighbors) - switch change.Position { - case "front": - destination = 0 - case "before", "after": - if anchor < 0 { - return 0, false, storeerr.ErrStateTransitionConflict - } - destination = anchor - if change.Position == "after" { - destination++ - } - } - var lower int64 - if destination > 0 { - lower = neighbors[destination-1] - } - if destination < len(neighbors) { - upper := neighbors[destination] - return lower + (upper-lower)/2, upper-lower > 1, nil - } - if lower > math.MaxInt64-1024 { - return 0, false, errors.New("backlog rank overflow") - } - return lower + 1024, true, nil -} diff --git a/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go b/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go index 0f0b5208a..5fbddde08 100644 --- a/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go +++ b/internal/storage/internal/agentexecution/hot_path_plans_integration_test.go @@ -155,7 +155,6 @@ model_call_contexts; ANALYZE model_outputs name, path string args []any }{ - {"ListExecutionBacklog", "internal/executiondb/commands_inputs.sql.go", []any{agent}}, {"ListContextEvents", "../dbsqlc/agent_turns.sql.go", []any{f.project, agent, int64(0), int64(100), int32(100)}}, {"CreateExecutionAttempt", "internal/executiondb/commands_models.sql.go", []any{uuid.New(), snapshot.Head.CurrentTurnID, "normal", int32(1), f.config, f.revision, int64(100), nil, lease, @@ -183,14 +182,7 @@ model_call_contexts; ANALYZE model_outputs t.Logf("%s iteration=%d plan=%s", tc.name, iteration, raw) require.NoError(t, nested.Rollback(t.Context())) } - if tc.name == "ListExecutionBacklog" { - results, err := tx.Conn().PgConn().Exec(t.Context(), "EXPLAIN (GENERIC_PLAN, FORMAT JSON) "+query).ReadAll() - require.NoError(t, err) - require.Len(t, results, 1) - require.Contains(t, string(results[0].Rows[0][0]), "agent_inputs_received_content_idx") - require.LessOrEqual(t, planBuffers(t, raw), 10) - require.Contains(t, string(raw), "agent_inputs_received_content_idx") - } + }) } } diff --git a/internal/storage/internal/agentexecution/internal/executiondb/batch.go b/internal/storage/internal/agentexecution/internal/executiondb/batch.go index 067f51e30..399ac3b78 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/batch.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/batch.go @@ -159,10 +159,12 @@ WHERE c.agent_id=$1 AND c.model_output_id=$2 AND NOT b.incomplete ), scope AS MATERIALIZED ( SELECT a.id,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now,t.id AS turn_id, coalesce(e.sequence,0)::bigint AS semantic_sequence,coalesce(e.created_at,'epoch'::timestamptz)::timestamptz AS semantic_time, - EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' - AND i.delivery_mode='steering' AND i.state='received') AS steering, - EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' - AND i.delivery_mode='queued' AND i.state='received') AS queued, + ((SELECT i.id FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='steering' AND i.state='received' + ORDER BY i.input_rank,i.queued_at,i.id LIMIT 1) IS NOT NULL)::boolean AS steering, + ((SELECT i.id FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='queued' AND i.state='received' + ORDER BY i.input_rank,i.queued_at,i.id LIMIT 1) IS NOT NULL)::boolean AS queued, coalesce(o.first_opening_sequence,0)::bigint AS first_opening_sequence, coalesce(o.last_opening_sequence,0)::bigint AS last_opening_sequence, coalesce(o.first_content_sequence,0)::bigint AS first_content_sequence diff --git a/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go b/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go index a0e90e27b..9a6489718 100644 --- a/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go +++ b/internal/storage/internal/agentexecution/internal/executiondb/commands_inputs.sql.go @@ -116,41 +116,63 @@ func (q *Queries) FindExecutionInput(ctx context.Context, db DBTX, arg FindExecu return i, err } -const listExecutionBacklog = `-- name: ListExecutionBacklog :many -SELECT id,delivery_mode,input_rank,queued_at,integration_target_id,actor_id FROM agent_inputs +const firstExecutionBacklogRank = `-- name: FirstExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs WHERE agent_id=$1 AND input_kind='content' AND state='received' -ORDER BY input_rank,queued_at,id + AND delivery_mode=$2 AND id<>$3 +ORDER BY input_rank,queued_at,id LIMIT 1 ` -type ListExecutionBacklogParams struct { +type FirstExecutionBacklogRankParams struct { AgentID uuid.UUID + Mode string + SkipID uuid.UUID } -type ListExecutionBacklogRow struct { - ID uuid.UUID - DeliveryMode string - InputRank int64 - QueuedAt time.Time - IntegrationTargetID *uuid.UUID - ActorID *uuid.UUID +func (q *Queries) FirstExecutionBacklogRank(ctx context.Context, db DBTX, arg FirstExecutionBacklogRankParams) (int64, error) { + row := db.QueryRow(ctx, firstExecutionBacklogRank, arg.AgentID, arg.Mode, arg.SkipID) + var input_rank int64 + err := row.Scan(&input_rank) + return input_rank, err } -func (q *Queries) ListExecutionBacklog(ctx context.Context, db DBTX, arg ListExecutionBacklogParams) ([]ListExecutionBacklogRow, error) { - rows, err := db.Query(ctx, listExecutionBacklog, arg.AgentID) +const getExecutionBacklogInputs = `-- name: GetExecutionBacklogInputs :many +SELECT id,state,input_kind,delivery_mode,input_rank,queued_at,admitted_event_id FROM agent_inputs +WHERE agent_id=$1 AND id=ANY($2::uuid[]) +` + +type GetExecutionBacklogInputsParams struct { + AgentID uuid.UUID + Ids []uuid.UUID +} + +type GetExecutionBacklogInputsRow struct { + ID uuid.UUID + State string + InputKind string + DeliveryMode string + InputRank int64 + QueuedAt time.Time + AdmittedEventID *uuid.UUID +} + +func (q *Queries) GetExecutionBacklogInputs(ctx context.Context, db DBTX, arg GetExecutionBacklogInputsParams) ([]GetExecutionBacklogInputsRow, error) { + rows, err := db.Query(ctx, getExecutionBacklogInputs, arg.AgentID, arg.Ids) if err != nil { return nil, err } defer rows.Close() - items := []ListExecutionBacklogRow{} + items := []GetExecutionBacklogInputsRow{} for rows.Next() { - var i ListExecutionBacklogRow + var i GetExecutionBacklogInputsRow if err := rows.Scan( &i.ID, + &i.State, + &i.InputKind, &i.DeliveryMode, &i.InputRank, &i.QueuedAt, - &i.IntegrationTargetID, - &i.ActorID, + &i.AdmittedEventID, ); err != nil { return nil, err } @@ -162,12 +184,120 @@ func (q *Queries) ListExecutionBacklog(ctx context.Context, db DBTX, arg ListExe return items, nil } +const lastExecutionBacklogRank = `-- name: LastExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs +WHERE agent_id=$1 AND input_kind='content' AND state='received' + AND delivery_mode=$2 AND id<>$3 +ORDER BY input_rank DESC,queued_at DESC,id DESC LIMIT 1 +` + +type LastExecutionBacklogRankParams struct { + AgentID uuid.UUID + Mode string + SkipID uuid.UUID +} + +func (q *Queries) LastExecutionBacklogRank(ctx context.Context, db DBTX, arg LastExecutionBacklogRankParams) (int64, error) { + row := db.QueryRow(ctx, lastExecutionBacklogRank, arg.AgentID, arg.Mode, arg.SkipID) + var input_rank int64 + err := row.Scan(&input_rank) + return input_rank, err +} + +const nextExecutionBacklogRank = `-- name: NextExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs +WHERE agent_id=$1 AND input_kind='content' AND state='received' + AND delivery_mode=$2 AND id<>$3 + AND (input_rank,queued_at,id)>($4::bigint,$5::timestamptz,$6::uuid) +ORDER BY input_rank,queued_at,id LIMIT 1 +` + +type NextExecutionBacklogRankParams struct { + AgentID uuid.UUID + Mode string + SkipID uuid.UUID + Rank int64 + QueuedAt time.Time + AnchorID uuid.UUID +} + +func (q *Queries) NextExecutionBacklogRank(ctx context.Context, db DBTX, arg NextExecutionBacklogRankParams) (int64, error) { + row := db.QueryRow(ctx, nextExecutionBacklogRank, + arg.AgentID, + arg.Mode, + arg.SkipID, + arg.Rank, + arg.QueuedAt, + arg.AnchorID, + ) + var input_rank int64 + err := row.Scan(&input_rank) + return input_rank, err +} + +const previousExecutionBacklogRank = `-- name: PreviousExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs +WHERE agent_id=$1 AND input_kind='content' AND state='received' + AND delivery_mode=$2 AND id<>$3 + AND (input_rank,queued_at,id)<($4::bigint,$5::timestamptz,$6::uuid) +ORDER BY input_rank DESC,queued_at DESC,id DESC LIMIT 1 +` + +type PreviousExecutionBacklogRankParams struct { + AgentID uuid.UUID + Mode string + SkipID uuid.UUID + Rank int64 + QueuedAt time.Time + AnchorID uuid.UUID +} + +func (q *Queries) PreviousExecutionBacklogRank(ctx context.Context, db DBTX, arg PreviousExecutionBacklogRankParams) (int64, error) { + row := db.QueryRow(ctx, previousExecutionBacklogRank, + arg.AgentID, + arg.Mode, + arg.SkipID, + arg.Rank, + arg.QueuedAt, + arg.AnchorID, + ) + var input_rank int64 + err := row.Scan(&input_rank) + return input_rank, err +} + +const rebalanceExecutionBacklog = `-- name: RebalanceExecutionBacklog :execrows +WITH ranked AS ( + SELECT pending.id,row_number() OVER (ORDER BY pending.input_rank,pending.queued_at,pending.id)*1024 AS rank + FROM agent_inputs pending WHERE pending.agent_id=$1 + AND pending.input_kind='content' AND pending.state='received' AND pending.delivery_mode=$2 +) +UPDATE agent_inputs i SET input_rank=r.rank FROM ranked r +WHERE i.agent_id=$1 AND i.id=r.id + AND i.input_kind='content' AND i.state='received' AND i.delivery_mode=$2 +` + +type RebalanceExecutionBacklogParams struct { + AgentID uuid.UUID + Mode string +} + +func (q *Queries) RebalanceExecutionBacklog(ctx context.Context, db DBTX, arg RebalanceExecutionBacklogParams) (int64, error) { + result, err := db.Exec(ctx, rebalanceExecutionBacklog, arg.AgentID, arg.Mode) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + const receiveExecutionInput = `-- name: ReceiveExecutionInput :one INSERT INTO agent_inputs(id,project_id,agent_id,state,input_rank,actor_id,input_kind,delivery_mode, integration_target_id,agent_config_id,idempotency_scope,input_idempotency_key,queued_at,metadata) -SELECT $1,a.project_id,a.id,'received',coalesce((SELECT max(i.input_rank)+1024 - FROM agent_inputs i WHERE i.agent_id=a.id AND i.state='received' AND i.delivery_mode=$2),1024), - $3,$4,$2,$5,$6, +SELECT $1,a.project_id,a.id,'received',CASE WHEN $2::text='content' THEN coalesce(( + SELECT i.input_rank+1024 FROM agent_inputs i WHERE i.agent_id=a.id AND i.state='received' + AND i.input_kind='content' AND i.delivery_mode=$3 + ORDER BY i.input_rank DESC,i.queued_at DESC,i.id DESC LIMIT 1),1024) ELSE 1024 END, + $4,$2,$3,$5,$6, $7,$8,statement_timestamp(),$9::jsonb FROM agents a WHERE a.id=$10 AND a.project_id=$11 AND a.state='active' RETURNING id,queued_at @@ -175,9 +305,9 @@ RETURNING id,queued_at type ReceiveExecutionInputParams struct { ID uuid.UUID + Kind string Mode string ActorID *uuid.UUID - Kind string TargetID *uuid.UUID ConfigID *uuid.UUID Scope *string @@ -195,9 +325,9 @@ type ReceiveExecutionInputRow struct { func (q *Queries) ReceiveExecutionInput(ctx context.Context, db DBTX, arg ReceiveExecutionInputParams) (ReceiveExecutionInputRow, error) { row := db.QueryRow(ctx, receiveExecutionInput, arg.ID, + arg.Kind, arg.Mode, arg.ActorID, - arg.Kind, arg.TargetID, arg.ConfigID, arg.Scope, @@ -240,3 +370,61 @@ func (q *Queries) ResolveExecutionInput(ctx context.Context, db DBTX, arg Resolv } return result.RowsAffected(), nil } + +const selectExecutionAdmission = `-- name: SelectExecutionAdmission :many +WITH selected AS MATERIALIZED ( + SELECT id,actor_id,integration_target_id,input_rank,queued_at FROM agent_inputs + WHERE agent_id=$2 AND input_kind='content' AND state='received' + AND delivery_mode=$1 + ORDER BY input_rank,queued_at,id LIMIT $3::bigint +) +SELECT s.id,s.actor_id,s.integration_target_id, + $1::text='queued' AND (SELECT i.id FROM agent_inputs i + WHERE i.agent_id=$2 AND i.input_kind='content' AND i.state='received' + AND i.delivery_mode='steering' ORDER BY i.input_rank,i.queued_at,i.id LIMIT 1) IS NOT NULL AS remaining_steering, + EXISTS(SELECT 1 FROM (SELECT i.id FROM agent_inputs i WHERE i.agent_id=$2 + AND i.input_kind='content' AND i.state='received' AND i.delivery_mode='queued' + ORDER BY i.input_rank DESC,i.queued_at DESC,i.id DESC LIMIT 1) pending + WHERE pending.id IS DISTINCT FROM (SELECT id FROM selected WHERE $1::text='queued' LIMIT 1)) AS remaining_queued +FROM selected s ORDER BY s.input_rank,s.queued_at,s.id +` + +type SelectExecutionAdmissionParams struct { + Mode string + AgentID uuid.UUID + RowLimit *int64 +} + +type SelectExecutionAdmissionRow struct { + ID uuid.UUID + ActorID *uuid.UUID + IntegrationTargetID *uuid.UUID + RemainingSteering *bool + RemainingQueued bool +} + +func (q *Queries) SelectExecutionAdmission(ctx context.Context, db DBTX, arg SelectExecutionAdmissionParams) ([]SelectExecutionAdmissionRow, error) { + rows, err := db.Query(ctx, selectExecutionAdmission, arg.Mode, arg.AgentID, arg.RowLimit) + if err != nil { + return nil, err + } + defer rows.Close() + items := []SelectExecutionAdmissionRow{} + for rows.Next() { + var i SelectExecutionAdmissionRow + if err := rows.Scan( + &i.ID, + &i.ActorID, + &i.IntegrationTargetID, + &i.RemainingSteering, + &i.RemainingQueued, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/storage/internal/agentexecution/load_count_integration_test.go b/internal/storage/internal/agentexecution/load_count_integration_test.go index dc62683b0..7c84576d7 100644 --- a/internal/storage/internal/agentexecution/load_count_integration_test.go +++ b/internal/storage/internal/agentexecution/load_count_integration_test.go @@ -17,17 +17,23 @@ type countingLoadTx struct { queries int facts []*pgx.QueuedQuery factLoads int + trace []ExecutionStatement } func (tx *countingLoadTx) Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error) { tx.statements++ tx.queries++ - return tx.Tx.Exec(ctx, sql, args...) + tag, err := tx.Tx.Exec(ctx, sql, args...) + tx.trace = append(tx.trace, ExecutionStatement{ + Query: pgx.QueuedQuery{SQL: sql, Arguments: args}, Affected: tag.RowsAffected(), + }) + return tag, err } func (tx *countingLoadTx) Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) { tx.statements++ tx.queries++ + tx.trace = append(tx.trace, ExecutionStatement{Query: pgx.QueuedQuery{SQL: sql, Arguments: args}}) rows, err := tx.Tx.Query(ctx, sql, args...) if err != nil { return nil, err @@ -42,6 +48,7 @@ func (tx *countingLoadTx) Query(ctx context.Context, sql string, args ...any) (p func (tx *countingLoadTx) QueryRow(ctx context.Context, sql string, args ...any) pgx.Row { tx.statements++ tx.queries++ + tx.trace = append(tx.trace, ExecutionStatement{Query: pgx.QueuedQuery{SQL: sql, Arguments: args}}) return tx.Tx.QueryRow(ctx, sql, args...) } @@ -56,6 +63,7 @@ func CountOrdinaryLoad(ctx context.Context, h *Handle) (ExecutionSnapshot, int, func (tx *countingLoadTx) SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults { tx.statements += batch.Len() for _, query := range batch.QueuedQueries { + tx.trace = append(tx.trace, ExecutionStatement{Query: *query}) if strings.HasPrefix(query.SQL, "-- name: LoadExecution") { tx.facts = append(tx.facts, query) } @@ -130,3 +138,16 @@ func CountExecutionTransition(ctx context.Context, u *Unit, run func()) (int, in err := u.Commit(ctx, "count transition") return tx.statements, tx.factLoads, err } + +type ExecutionStatement struct { + Query pgx.QueuedQuery + Affected int64 +} + +func TraceExecutionCommand(u *Unit, run func() error) ([]ExecutionStatement, error) { + tx := &countingLoadTx{Tx: u.tx} + u.tx = tx + defer func() { u.tx = tx.Tx }() + err := run() + return tx.trace, err +} diff --git a/internal/storage/internal/agentexecution/queries/commands_inputs.sql b/internal/storage/internal/agentexecution/queries/commands_inputs.sql index f9b883f19..61f1d7a92 100644 --- a/internal/storage/internal/agentexecution/queries/commands_inputs.sql +++ b/internal/storage/internal/agentexecution/queries/commands_inputs.sql @@ -7,17 +7,71 @@ FROM agent_inputs WHERE agent_id=sqlc.arg(agent_id) AND -- name: ReceiveExecutionInput :one INSERT INTO agent_inputs(id,project_id,agent_id,state,input_rank,actor_id,input_kind,delivery_mode, integration_target_id,agent_config_id,idempotency_scope,input_idempotency_key,queued_at,metadata) -SELECT sqlc.arg(id),a.project_id,a.id,'received',coalesce((SELECT max(i.input_rank)+1024 - FROM agent_inputs i WHERE i.agent_id=a.id AND i.state='received' AND i.delivery_mode=sqlc.arg(mode)),1024), +SELECT sqlc.arg(id),a.project_id,a.id,'received',CASE WHEN sqlc.arg(kind)::text='content' THEN coalesce(( + SELECT i.input_rank+1024 FROM agent_inputs i WHERE i.agent_id=a.id AND i.state='received' + AND i.input_kind='content' AND i.delivery_mode=sqlc.arg(mode) + ORDER BY i.input_rank DESC,i.queued_at DESC,i.id DESC LIMIT 1),1024) ELSE 1024 END, sqlc.narg(actor_id),sqlc.arg(kind),sqlc.arg(mode),sqlc.narg(target_id),sqlc.narg(config_id), sqlc.narg(scope),sqlc.narg(key),statement_timestamp(),sqlc.arg(metadata)::jsonb FROM agents a WHERE a.id=sqlc.arg(agent_id) AND a.project_id=sqlc.arg(project_id) AND a.state='active' RETURNING id,queued_at; --- name: ListExecutionBacklog :many -SELECT id,delivery_mode,input_rank,queued_at,integration_target_id,actor_id FROM agent_inputs +-- name: SelectExecutionAdmission :many +WITH selected AS MATERIALIZED ( + SELECT id,actor_id,integration_target_id,input_rank,queued_at FROM agent_inputs + WHERE agent_id=sqlc.arg(agent_id) AND input_kind='content' AND state='received' + AND delivery_mode=sqlc.arg(mode) + ORDER BY input_rank,queued_at,id LIMIT sqlc.narg(row_limit)::bigint +) +SELECT s.id,s.actor_id,s.integration_target_id, + sqlc.arg(mode)::text='queued' AND (SELECT i.id FROM agent_inputs i + WHERE i.agent_id=sqlc.arg(agent_id) AND i.input_kind='content' AND i.state='received' + AND i.delivery_mode='steering' ORDER BY i.input_rank,i.queued_at,i.id LIMIT 1) IS NOT NULL AS remaining_steering, + EXISTS(SELECT 1 FROM (SELECT i.id FROM agent_inputs i WHERE i.agent_id=sqlc.arg(agent_id) + AND i.input_kind='content' AND i.state='received' AND i.delivery_mode='queued' + ORDER BY i.input_rank DESC,i.queued_at DESC,i.id DESC LIMIT 1) pending + WHERE pending.id IS DISTINCT FROM (SELECT id FROM selected WHERE sqlc.arg(mode)::text='queued' LIMIT 1)) AS remaining_queued +FROM selected s ORDER BY s.input_rank,s.queued_at,s.id; + +-- name: GetExecutionBacklogInputs :many +SELECT id,state,input_kind,delivery_mode,input_rank,queued_at,admitted_event_id FROM agent_inputs +WHERE agent_id=sqlc.arg(agent_id) AND id=ANY(sqlc.arg(ids)::uuid[]); + +-- name: FirstExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs WHERE agent_id=sqlc.arg(agent_id) AND input_kind='content' AND state='received' -ORDER BY input_rank,queued_at,id; + AND delivery_mode=sqlc.arg(mode) AND id<>sqlc.arg(skip_id) +ORDER BY input_rank,queued_at,id LIMIT 1; + +-- name: LastExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs +WHERE agent_id=sqlc.arg(agent_id) AND input_kind='content' AND state='received' + AND delivery_mode=sqlc.arg(mode) AND id<>sqlc.arg(skip_id) +ORDER BY input_rank DESC,queued_at DESC,id DESC LIMIT 1; + +-- name: PreviousExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs +WHERE agent_id=sqlc.arg(agent_id) AND input_kind='content' AND state='received' + AND delivery_mode=sqlc.arg(mode) AND id<>sqlc.arg(skip_id) + AND (input_rank,queued_at,id)<(sqlc.arg(rank)::bigint,sqlc.arg(queued_at)::timestamptz,sqlc.arg(anchor_id)::uuid) +ORDER BY input_rank DESC,queued_at DESC,id DESC LIMIT 1; + +-- name: NextExecutionBacklogRank :one +SELECT input_rank FROM agent_inputs +WHERE agent_id=sqlc.arg(agent_id) AND input_kind='content' AND state='received' + AND delivery_mode=sqlc.arg(mode) AND id<>sqlc.arg(skip_id) + AND (input_rank,queued_at,id)>(sqlc.arg(rank)::bigint,sqlc.arg(queued_at)::timestamptz,sqlc.arg(anchor_id)::uuid) +ORDER BY input_rank,queued_at,id LIMIT 1; + +-- name: RebalanceExecutionBacklog :execrows +WITH ranked AS ( + SELECT pending.id,row_number() OVER (ORDER BY pending.input_rank,pending.queued_at,pending.id)*1024 AS rank + FROM agent_inputs pending WHERE pending.agent_id=sqlc.arg(agent_id) + AND pending.input_kind='content' AND pending.state='received' AND pending.delivery_mode=sqlc.arg(mode) +) +UPDATE agent_inputs i SET input_rank=r.rank FROM ranked r +WHERE i.agent_id=sqlc.arg(agent_id) AND i.id=r.id + AND i.input_kind='content' AND i.state='received' AND i.delivery_mode=sqlc.arg(mode); -- name: ChangeExecutionInput :execrows UPDATE agent_inputs SET delivery_mode=sqlc.arg(mode),input_rank=sqlc.arg(rank), diff --git a/internal/storage/internal/agentexecution/queries/load.sql b/internal/storage/internal/agentexecution/queries/load.sql index 70c0183cf..69e24bcae 100644 --- a/internal/storage/internal/agentexecution/queries/load.sql +++ b/internal/storage/internal/agentexecution/queries/load.sql @@ -30,10 +30,12 @@ WHERE c.agent_id=sqlc.arg(agent_id) AND c.model_output_id=sqlc.narg(tool_output_ ), scope AS MATERIALIZED ( SELECT a.id,a.root_agent_id,a.state,statement_timestamp()::timestamptz AS database_now,t.id AS turn_id, coalesce(e.sequence,0)::bigint AS semantic_sequence,coalesce(e.created_at,'epoch'::timestamptz)::timestamptz AS semantic_time, - EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' - AND i.delivery_mode='steering' AND i.state='received') AS steering, - EXISTS(SELECT 1 FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' - AND i.delivery_mode='queued' AND i.state='received') AS queued, + ((SELECT i.id FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='steering' AND i.state='received' + ORDER BY i.input_rank,i.queued_at,i.id LIMIT 1) IS NOT NULL)::boolean AS steering, + ((SELECT i.id FROM agent_inputs i WHERE i.agent_id=a.id AND i.input_kind='content' + AND i.delivery_mode='queued' AND i.state='received' + ORDER BY i.input_rank,i.queued_at,i.id LIMIT 1) IS NOT NULL)::boolean AS queued, coalesce(o.first_opening_sequence,0)::bigint AS first_opening_sequence, coalesce(o.last_opening_sequence,0)::bigint AS last_opening_sequence, coalesce(o.first_content_sequence,0)::bigint AS first_content_sequence diff --git a/migrations/000052_agent_execution_state.sql b/migrations/000052_agent_execution_state.sql index 3a6f1a7e6..2f8fe74bb 100644 --- a/migrations/000052_agent_execution_state.sql +++ b/migrations/000052_agent_execution_state.sql @@ -456,9 +456,12 @@ CREATE INDEX tool_calls_runtime_reference_idx DROP INDEX tool_calls_runtime_lock_idx; CREATE INDEX agent_inputs_received_content_idx - ON agent_inputs(agent_id, input_rank, queued_at, id) + ON agent_inputs(agent_id, delivery_mode, input_rank, queued_at, id) WHERE state = 'received' AND input_kind = 'content'; +DROP INDEX agent_inputs_content_queued_idx; +DROP INDEX agent_inputs_content_steering_idx; + CREATE UNIQUE INDEX agents_org_project_id_key ON agents(org_id, project_id, id); From 06ef15625bafece3319a4e88344fffd1a950875d Mon Sep 17 00:00:00 2001 From: Kartik Sarangmath <22248603+ksarangmath@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:41:13 -0700 Subject: [PATCH 09/19] feat: cap pending content inputs at 50 per agent - Reject a new content submission with agent_backlog_full (HTTP 409) when the agent already has 50 pending inputs, counting queued and steering together. The check runs under the agent lock after the idempotency lookup, so replays still succeed, and counts at most 50 index entries. - Accept a child's results, failures and questions to its parent past the cap; they still count toward it. - Leave integration inbox deliveries pending for retry, record a skipped cron occurrence, and fail a parent's message to a full subagent with a tool error the parent can read. - Show the error in the composer and restore the exact unsent draft without discarding text typed while the send was in flight. Co-Authored-By: Claude Opus 5.5 (1M context) --- api/openapi/openapi.yaml | 3 + docs/api-reference/openapi.yaml | 3 + docs/api/overview.mdx | 2 +- docs/events/sending-input.mdx | 2 + .../components/agents/AgentComposer.test.tsx | 161 ++ .../src/components/agents/AgentComposer.tsx | 15 +- .../src/domains/agent-chat-session.test.ts | 18 + .../packages/sdk/src/generated/types.gen.ts | 2 +- .../packages/sdk/src/generated/zod.gen.ts | 1 + internal/crontrigger/service.go | 12 +- .../dispatch_transaction_integration_test.go | 46 + internal/harness/tools/subagent_tools.go | 10 +- internal/httpapi/apierror/errors.go | 2 + internal/httpapi/apierror/errors_test.go | 8 + .../backlog_capacity_integration_test.go | 36 + internal/httpapi/openapi/api.gen.go | 1894 +++++++++-------- .../integration/worker_integration_test.go | 68 + .../cron_trigger_service_integration_test.go | 41 + .../backlog_capacity_integration_test.go | 165 ++ .../backlog_cost_integration_test.go | 34 +- .../agentexecution/commands_agents.go | 2 +- .../commands_family_integration_test.go | 53 +- .../agentexecution/commands_inputs.go | 27 +- .../executiondb/commands_inputs.sql.go | 19 + .../queries/commands_inputs.sql | 7 + .../transition_count_integration_test.go | 2 +- internal/storage/storeerr/errors.go | 4 + 27 files changed, 1660 insertions(+), 977 deletions(-) create mode 100644 frontend/apps/web/src/components/agents/AgentComposer.test.tsx create mode 100644 internal/httpapi/backlog_capacity_integration_test.go create mode 100644 internal/storage/internal/agentexecution/backlog_capacity_integration_test.go diff --git a/api/openapi/openapi.yaml b/api/openapi/openapi.yaml index 7799b8c4c..bcec0e6f2 100644 --- a/api/openapi/openapi.yaml +++ b/api/openapi/openapi.yaml @@ -8476,6 +8476,7 @@ components: - internal_error - service_unavailable - idempotency_key_conflict + - agent_backlog_full - state_transition_conflict - managed_work_admission_denied - pending_work @@ -8506,6 +8507,7 @@ components: - internal_error - service_unavailable - idempotency_key_conflict + - agent_backlog_full - state_transition_conflict - managed_work_admission_denied - pending_work @@ -8530,6 +8532,7 @@ components: - expected_digest_required - not_a_file - idempotency_key_conflict + - agent_backlog_full - state_transition_conflict - pending_work - not_wake_capable diff --git a/docs/api-reference/openapi.yaml b/docs/api-reference/openapi.yaml index 37fc3b35a..1fd5a3aac 100644 --- a/docs/api-reference/openapi.yaml +++ b/docs/api-reference/openapi.yaml @@ -8477,6 +8477,7 @@ components: - internal_error - service_unavailable - idempotency_key_conflict + - agent_backlog_full - state_transition_conflict - managed_work_admission_denied - pending_work @@ -8507,6 +8508,7 @@ components: - internal_error - service_unavailable - idempotency_key_conflict + - agent_backlog_full - state_transition_conflict - managed_work_admission_denied - pending_work @@ -8531,6 +8533,7 @@ components: - expected_digest_required - not_a_file - idempotency_key_conflict + - agent_backlog_full - state_transition_conflict - pending_work - not_wake_capable diff --git a/docs/api/overview.mdx b/docs/api/overview.mdx index 26411dd23..23e55905d 100644 --- a/docs/api/overview.mdx +++ b/docs/api/overview.mdx @@ -84,7 +84,7 @@ Branch on `code` — it's a stable contract, enumerated in the OpenAPI spec, and Two rules worth internalizing: - **Resources you cannot read return `404`, not `403`.** The API doesn't confirm the existence of things outside your access. A `403` always means "the resource is visible to you, but your role doesn't permit this action." -- **Retry `424`, `429` and `5xx` with backoff, reusing the same `Idempotency-Key`** so a retry can never act twice. `424` with code `upstream_unavailable` means a server you configured, such as an MCP server, did not respond or failed; retries succeed once that server recovers. A `409` from MCP tool discovery means another refresh of the same server or auth secret is in progress; retry it shortly. Other `4xx` errors won't change on retry — fix the request instead. +- **Retry `424`, `429` and `5xx` with backoff, reusing the same `Idempotency-Key`** so a retry can never act twice. `424` with code `upstream_unavailable` means a server you configured, such as an MCP server, did not respond or failed; retries succeed once that server recovers. A `409` from MCP tool discovery means another refresh of the same server or auth secret is in progress; retry it shortly. A `409` with code `agent_backlog_full` means the agent has 50 pending messages; wait for processing or cancel a pending message before retrying. Other `4xx` errors won't change on retry — fix the request instead. ## Request size limits diff --git a/docs/events/sending-input.mdx b/docs/events/sending-input.mdx index 08195f69f..b2e854bb1 100644 --- a/docs/events/sending-input.mdx +++ b/docs/events/sending-input.mdx @@ -29,6 +29,8 @@ export AGENT="agt_v4qtrwz3jehcyd5n6a2bfgik7m" ## Send a message +An agent accepts up to 50 pending messages across queued and steering inputs; further submissions return `409` with code `agent_backlog_full` until messages are processed or canceled, while child lifecycle reports are always delivered. + The examples assume the client, `orgID`, and `projectID` setup from the [quickstart](/quickstart), with the `$AGENT` ID above available as `agentID` for the SDK. The CLI has no command for sending input — use the REST API or the SDK. diff --git a/frontend/apps/web/src/components/agents/AgentComposer.test.tsx b/frontend/apps/web/src/components/agents/AgentComposer.test.tsx new file mode 100644 index 000000000..72ac08e73 --- /dev/null +++ b/frontend/apps/web/src/components/agents/AgentComposer.test.tsx @@ -0,0 +1,161 @@ +/** @vitest-environment happy-dom */ + +import type { UseAgentChatResult } from '@omnara/react' +import { ApiError } from '@omnara/sdk' +import { act, useState } from 'react' +import { createRoot } from 'react-dom/client' +import { expect, it, vi } from 'vitest' + +import { MessageScrollerProvider } from '@/components/ui/message-scroller' +import { enableReactActEnvironment } from '@/test/react-act' + +import { AgentComposer } from './AgentComposer' + +const message = + 'This agent already has 50 pending messages. Wait for messages to be processed or cancel a pending message, then try again.' + +type Send = UseAgentChatResult['sendMessage'] + +interface InFlightSend { + settle?: (failure?: Error) => void +} + +function Composer({ send }: { send: Send }) { + const [error, setError] = useState() + const chat: UseAgentChatResult = { + events: [], + messages: [], + status: 'ready', + isWorking: false, + error, + streamError: undefined, + historyStatus: 'success', + historyError: null, + retryHistory: vi.fn(), + reconnect: vi.fn(), + hasOlderMessages: false, + isLoadingOlderMessages: false, + loadOlderMessages: vi.fn(), + inputBacklog: { + inputs: [], + actionPending: false, + beginCancellation: vi.fn(), + cancel: vi.fn(), + promote: vi.fn(), + move: vi.fn(), + }, + sendMessage: async (input) => { + setError(undefined) + try { + await send(input) + } catch (cause) { + setError(cause instanceof Error ? cause : new Error(String(cause))) + throw cause + } + }, + } + return ( +