test: pin MultiXact-free fused acceptance and the ready-index claim plan

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: ksarangmath

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) {}

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)
+}