diff --git a/internal/storage/internal/agentexecution/hot_path_migration_integration_test.go b/internal/storage/internal/agentexecution/hot_path_migration_integration_test.go new file mode 100644 index 000000000..5d37cd482 --- /dev/null +++ b/internal/storage/internal/agentexecution/hot_path_migration_integration_test.go @@ -0,0 +1,108 @@ +//go:build integration + +package agentexecution_test + +import ( + "os" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestHotPathMigrationConcurrentBuildAndValidation(t *testing.T) { + f := executionFixtureThrough(t, 53) + ids := make([]uuid.UUID, 50000) + for i := range ids { + ids[i] = uuid.New() + } + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + _, err = tx.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 = tx.Exec(t.Context(), ` +INSERT INTO agent_inputs(id,project_id,agent_id,state,input_kind,delivery_mode,queued_at) + SELECT id,$2,id,'received','content','queued',now() FROM unnest($1::uuid[]) id +`, ids, f.project) + 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, +created_at,idempotency_key) + SELECT id,id,id,1,'agent_input',id,true,now(),'agent_input:'||id FROM unnest($1::uuid[]) id +`, ids) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +INSERT INTO agent_turns(id,agent_id,turn_sequence,latest_event_id,latest_semantic_event_id) + SELECT id,id,1,id,id FROM unnest($1::uuid[]) id +`, ids) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), ` +UPDATE agent_inputs SET state='resolved',admitted_event_id=id,admitted_at=now(),resolved_at=now() +WHERE agent_id=ANY($1::uuid[]) +`, ids) + require.NoError(t, err) + _, err = tx.Exec(t.Context(), `ANALYZE agents; ANALYZE agent_inputs; ANALYZE agent_events; ANALYZE agent_turns`) + require.NoError(t, err) + _, 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 id,$2,$3,id,id,'normal',1,$4,$5,1,gen_random_uuid(),now(),ARRAY[id],1 FROM unnest($1::uuid[]) +id +`, ids, f.org, f.project, f.config, f.revision) + require.NoError(t, err) + require.NoError(t, tx.Commit(t.Context())) + migration, err := os.ReadFile("../../../../migrations/000054_hot_path_references.sql") + require.NoError(t, err) + require.Contains(t, string(migration), "-- +goose NO TRANSACTION") + statements := strings.Split(string(migration), ";") + for _, statement := range statements { + statement = strings.TrimSpace(statement) + if statement == "" { + continue + } + start := time.Now() + _, err = f.pool.Exec(t.Context(), statement) + require.NoError(t, err, statement) + t.Logf("rows=50000 duration=%s statement=%s", time.Since(start), statement) + } + for _, table := range []string{"model_call_contexts", "processes", "process_actions"} { + tx, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + _, err = tx.Exec(t.Context(), "ALTER TABLE "+table+" DROP CONSTRAINT "+table+"_agent_scope_fkey") + require.NoError(t, err) + _, err = tx.Exec(t.Context(), "ALTER TABLE "+table+" ADD CONSTRAINT "+table+ + "_agent_scope_fkey FOREIGN KEY(org_id,project_id,agent_id) REFERENCES agents(org_id,project_id,id) NOT VALID") + require.NoError(t, err) + require.NoError(t, tx.Commit(t.Context())) + tx, err = f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = tx.Rollback(t.Context()) }() + start := time.Now() + _, err = tx.Exec(t.Context(), "ALTER TABLE "+table+" VALIDATE CONSTRAINT "+table+"_agent_scope_fkey") + require.NoError(t, err) + var locks []string + require.NoError(t, tx.QueryRow(t.Context(), ` +SELECT array_agg(DISTINCT mode ORDER BY mode) FROM pg_locks + WHERE pid=pg_backend_pid() AND relation=ANY(ARRAY['agents'::regclass,$1::regclass]) +`, table).Scan(&locks)) + t.Logf("%s validation duration=%s locks=%v", table, time.Since(start), locks) + require.NotContains(t, locks, "AccessExclusiveLock") + require.NotContains(t, locks, "ShareRowExclusiveLock") + require.NotContains(t, locks, "ShareLock") + writer, err := f.pool.Begin(t.Context()) + require.NoError(t, err) + defer func() { _ = writer.Rollback(t.Context()) }() + _, err = writer.Exec(t.Context(), "LOCK TABLE "+table+" IN ROW EXCLUSIVE MODE NOWAIT") + require.NoError(t, err) + require.NoError(t, writer.Rollback(t.Context())) + require.NoError(t, tx.Commit(t.Context())) + } +} 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/internal/storage/internal/agentexecution/reconstruction_integration_test.go b/internal/storage/internal/agentexecution/reconstruction_integration_test.go index 485b2d625..25317c455 100644 --- a/internal/storage/internal/agentexecution/reconstruction_integration_test.go +++ b/internal/storage/internal/agentexecution/reconstruction_integration_test.go @@ -38,6 +38,11 @@ type executionFixture struct { } func releasedExecutionFixture(t *testing.T) executionFixture { + t.Helper() + return executionFixtureThrough(t, 51) +} + +func executionFixtureThrough(t *testing.T, lastVersion int) executionFixture { t.Helper() ctx := t.Context() pool := integrationdb.OpenUnmigratedPool(t, ctx) @@ -50,7 +55,7 @@ func releasedExecutionFixture(t *testing.T) executionFixture { } version, parseErr := strconv.Atoi(strings.Split(entry.Name(), "_")[0]) require.NoError(t, parseErr) - if version > 51 { + if version > lastVersion { continue } body, readErr := os.ReadFile(filepath.Join("../../../../migrations", entry.Name())) diff --git a/migrations/000054_hot_path_references.sql b/migrations/000054_hot_path_references.sql new file mode 100644 index 000000000..50ef9d385 --- /dev/null +++ b/migrations/000054_hot_path_references.sql @@ -0,0 +1,43 @@ +-- +goose NO TRANSACTION +-- +goose Up + +CREATE INDEX CONCURRENTLY tool_calls_runtime_reference_idx + ON tool_calls(runtime_lock_id) WHERE runtime_lock_id IS NOT NULL; + +DROP INDEX CONCURRENTLY tool_calls_runtime_lock_idx; + +CREATE INDEX CONCURRENTLY 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 CONCURRENTLY 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) NOT VALID; + +ALTER TABLE processes ADD CONSTRAINT processes_agent_scope_fkey + FOREIGN KEY (org_id, project_id, agent_id) REFERENCES agents(org_id, project_id, id) NOT VALID; + +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) NOT VALID; + +ALTER TABLE model_call_contexts VALIDATE CONSTRAINT model_call_contexts_agent_scope_fkey; +ALTER TABLE processes VALIDATE CONSTRAINT processes_agent_scope_fkey; +ALTER TABLE process_actions VALIDATE CONSTRAINT process_actions_agent_scope_fkey; + +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;