From 9400cbe5b2c50329c8ba40dacc0ca4412208b34c Mon Sep 17 00:00:00 2001 From: Delicious233 <101502465+DeliciousBuding@users.noreply.github.com> Date: Sun, 6 Sep 2026 12:39:21 +0800 Subject: [PATCH] fix(agent): propagate custom-agent request cancellation Bind every custom-agent CRUD database operation to the caller context. Cover canceled mutations and the ownership-read boundary in L0; prove HTTP deadline/disconnect releases blocked PostgreSQL queries and pool capacity before unlock in L1. Refs #2329 Refs #2256 Co-authored-by: Codex --- .../internal/service/agent/agent_custom.go | 12 +- .../agent/agent_custom_context_test.go | 111 +++++++++++ .../integration/custom_agent_context_test.go | 173 ++++++++++++++++++ 3 files changed, 290 insertions(+), 6 deletions(-) create mode 100644 hub-server/internal/service/agent/agent_custom_context_test.go create mode 100644 hub-server/tests/integration/custom_agent_context_test.go diff --git a/hub-server/internal/service/agent/agent_custom.go b/hub-server/internal/service/agent/agent_custom.go index 19a3f6fde..9ee7ed5b7 100644 --- a/hub-server/internal/service/agent/agent_custom.go +++ b/hub-server/internal/service/agent/agent_custom.go @@ -23,7 +23,7 @@ func (s *Service) CreateCustomAgent(ctx context.Context, ownerID, name, avatarUR ToolWhitelist: toolWhitelist, ModelParams: modelParams, } - if err := repository.CreateCustomAgent(s.db, ca); err != nil { + if err := repository.CreateCustomAgent(s.db.WithContext(ctx), ca); err != nil { return nil, err } return ca, nil @@ -31,7 +31,7 @@ func (s *Service) CreateCustomAgent(ctx context.Context, ownerID, name, avatarUR // ListCustomAgents returns all custom agents owned by the given user. func (s *Service) ListCustomAgents(ctx context.Context, ownerID string) ([]model.CustomAgent, error) { - return repository.ListCustomAgentsByOwner(s.db, ownerID) + return repository.ListCustomAgentsByOwner(s.db.WithContext(ctx), ownerID) } // UpdateCustomAgent updates an existing custom agent, verifying ownership. @@ -49,7 +49,7 @@ func (s *Service) ListCustomAgents(ctx context.Context, ownerID string) ([]model // tool_whitelist and model_params with omitempty, so an omitted value must not // be flattened to "" by a write that does include those columns. func (s *Service) UpdateCustomAgent(ctx context.Context, ownerID string, ca *model.CustomAgent) error { - existing, err := repository.GetCustomAgentByID(s.db, ca.ID) + existing, err := repository.GetCustomAgentByID(s.db.WithContext(ctx), ca.ID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return errcode.AgentNotFound @@ -72,12 +72,12 @@ func (s *Service) UpdateCustomAgent(ctx context.Context, ownerID string, ca *mod // repository.UpdateCustomAgent puts the not-deleted guard inside the UPDATE // and reports zero matched rows as ErrRecordNotFound, so that race surfaces // as the same 404 the read path returns instead of resurrecting the row. - return repository.WrapNotFound(repository.UpdateCustomAgent(s.db, ca), errcode.AgentNotFound) + return repository.WrapNotFound(repository.UpdateCustomAgent(s.db.WithContext(ctx), ca), errcode.AgentNotFound) } // DeleteCustomAgent soft-deletes a custom agent, verifying ownership. func (s *Service) DeleteCustomAgent(ctx context.Context, ownerID, id string) error { - ca, err := repository.GetCustomAgentByID(s.db, id) + ca, err := repository.GetCustomAgentByID(s.db.WithContext(ctx), id) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return errcode.AgentNotFound @@ -87,5 +87,5 @@ func (s *Service) DeleteCustomAgent(ctx context.Context, ownerID, id string) err if ca.OwnerUserID != ownerID { return errcode.AgentNotFound } - return repository.SoftDeleteCustomAgent(s.db, id) + return repository.SoftDeleteCustomAgent(s.db.WithContext(ctx), id) } diff --git a/hub-server/internal/service/agent/agent_custom_context_test.go b/hub-server/internal/service/agent/agent_custom_context_test.go new file mode 100644 index 000000000..5c3bd0516 --- /dev/null +++ b/hub-server/internal/service/agent/agent_custom_context_test.go @@ -0,0 +1,111 @@ +package agent + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + + "github.com/agenthub/hub-server/internal/model" +) + +const customAgentContextID = "context-agent" +const customAgentContextOwner = "context-owner" + +func customAgentContextFixture(t *testing.T) (*Service, *gorm.DB) { + t.Helper() + database := newCustomAgentUpdateDB(t) + sqlDB, err := database.DB() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, sqlDB.Close()) }) + seedCustomAgent(t, database, customAgentContextID, customAgentContextOwner, "Before cancellation", time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)) + return &Service{db: database}, database +} + +func assertCustomAgentContextUnchanged(t *testing.T, database *gorm.DB) { + t.Helper() + var count int64 + require.NoError(t, database.Model(&model.CustomAgent{}).Count(&count).Error) + assert.Equal(t, int64(1), count, "a canceled create must not insert a late row") + var row model.CustomAgent + require.NoError(t, database.Where("id = ?", customAgentContextID).First(&row).Error) + assert.Equal(t, "Before cancellation", row.Name) + assert.Equal(t, customAgentContextOwner, row.OwnerUserID) + assert.Nil(t, row.DeletedAt, "a canceled delete must not hide the agent") + require.NotNil(t, row.OutputSchema) + assert.JSONEq(t, customAgentUpdateTestSchema, string(*row.OutputSchema)) +} + +func updateCustomAgentWithContext(s *Service, ctx context.Context) error { + return s.UpdateCustomAgent(ctx, customAgentContextOwner, &model.CustomAgent{ + ID: customAgentContextID, Name: "After cancellation", AgentType: "claude-code", SystemPrompt: "Changed prompt", + }) +} + +func TestCustomAgentCanceledContext(t *testing.T) { + tests := []struct { + name string + call func(*Service, context.Context) error + }{ + {"create", func(s *Service, ctx context.Context) error { + _, err := s.CreateCustomAgent(ctx, customAgentContextOwner, "Unexpected agent", "", "claude-code", "Fixture prompt", "[]", "[]", "{}") + return err + }}, + {"list", func(s *Service, ctx context.Context) error { + _, err := s.ListCustomAgents(ctx, customAgentContextOwner) + return err + }}, + {"update", updateCustomAgentWithContext}, + {"delete", func(s *Service, ctx context.Context) error { + return s.DeleteCustomAgent(ctx, customAgentContextOwner, customAgentContextID) + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s, database := customAgentContextFixture(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + assert.ErrorIs(t, tt.call(s, ctx), context.Canceled) + assertCustomAgentContextUnchanged(t, database) + // Request binding must not poison the shared handle for the next caller. + agents, err := s.ListCustomAgents(context.Background(), customAgentContextOwner) + require.NoError(t, err) + assert.Len(t, agents, 1) + }) + } +} + +func TestCustomAgentCancelAfterOwnershipRead(t *testing.T) { + tests := []struct { + name string + call func(*Service, context.Context) error + }{ + {"update", updateCustomAgentWithContext}, + {"delete", func(s *Service, ctx context.Context) error { + return s.DeleteCustomAgent(ctx, customAgentContextOwner, customAgentContextID) + }}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s, database := customAgentContextFixture(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + // Cancel only after the real ownership SELECT has returned. An entry-only + // ctx.Err check cannot protect the subsequent UPDATE or soft DELETE. + const callback = "test:cancel-custom-agent-after-owner-read" + require.NoError(t, database.Callback().Query().After("gorm:query").Register(callback, func(tx *gorm.DB) { + if tx.Statement.Table == "custom_agents" && tx.Error == nil && tx.RowsAffected > 0 { + cancel() + } + })) + t.Cleanup(func() { require.NoError(t, database.Callback().Query().Remove(callback)) }) + err := tt.call(s, ctx) + require.ErrorIs(t, ctx.Err(), context.Canceled, "the cancellation interleaving must have happened") + assert.ErrorIs(t, err, context.Canceled) + assertCustomAgentContextUnchanged(t, database) + }) + } +} diff --git a/hub-server/tests/integration/custom_agent_context_test.go b/hub-server/tests/integration/custom_agent_context_test.go new file mode 100644 index 000000000..7c3daebf1 --- /dev/null +++ b/hub-server/tests/integration/custom_agent_context_test.go @@ -0,0 +1,173 @@ +//go:build integration + +package integration + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/postgres" + "gorm.io/gorm" + + "github.com/agenthub/hub-server/internal/config" + "github.com/agenthub/hub-server/internal/handler" + "github.com/agenthub/hub-server/internal/middleware" + "github.com/agenthub/hub-server/internal/service/agent" + "github.com/agenthub/pkg/testkit" +) + +// This is L1: real HTTP, handler, service and PostgreSQL; identity is a fixture. +// Canceling an HTTP client alone is insufficient evidence: the server-side SQL +// and the queued pool caller must finish BEFORE the blocking lock is released. +func TestCustomAgentRequestCancellationReleasesPool(t *testing.T) { + for _, mode := range []string{"deadline", "disconnect"} { + t.Run(mode, func(t *testing.T) { + t.Cleanup(func() { CleanDB(t, db) }) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + t.Cleanup(cancel) + u := register(t, "custom-context-"+mode, "pass1234", "Context fixture") + + // A private pool makes saturation observable without changing the + // shared integration server's pool or any runtime configuration. + dialect, ok := db.Dialector.(*postgres.Dialector) + require.True(t, ok) + scopedDB, err := gorm.Open(postgres.Open(dialect.DSN), &gorm.Config{ + Logger: db.Logger, SkipDefaultTransaction: db.SkipDefaultTransaction, + PrepareStmt: db.PrepareStmt, + }) + require.NoError(t, err) + pool, err := scopedDB.DB() + require.NoError(t, err) + pool.SetMaxOpenConns(1) + pool.SetMaxIdleConns(1) + t.Cleanup(func() { require.NoError(t, pool.Close()) }) + // Cancellation, not a server-side statement/lock timeout, must end SQL. + _, err = pool.ExecContext(ctx, "SET statement_timeout = 0") + require.NoError(t, err) + _, err = pool.ExecContext(ctx, "SET lock_timeout = 0") + require.NoError(t, err) + var queryPID int + require.NoError(t, pool.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&queryPID)) + + service := agent.NewService(scopedDB, eventBus, mgr, testCacheClient, nil, config.EdgeDispatchConfig{}, nil, "") + created, err := service.CreateCustomAgent(ctx, u.ID, "Still present", "", "codex", "Fixture prompt", "[]", "[]", "{}") + require.NoError(t, err) + + sharedPool, err := db.DB() + require.NoError(t, err) + lock, err := sharedPool.BeginTx(ctx, nil) + require.NoError(t, err) + unlock := sync.OnceFunc(func() { require.NoError(t, lock.Rollback()) }) + t.Cleanup(unlock) + _, err = lock.ExecContext(ctx, "LOCK TABLE custom_agents IN ACCESS EXCLUSIVE MODE") + require.NoError(t, err) + + deadline := time.Second + if mode == "disconnect" { + deadline = 10 * time.Second + } + started, serverDone := make(chan struct{}), make(chan struct{}) + var requestCtx context.Context + r := gin.New() + r.Use(func(c *gin.Context) { defer close(serverDone); c.Next() }) + r.Use(middleware.Timeout(deadline)) + h := handler.NewCustomAgentHandler(service) + r.GET("/web/custom-agents", func(c *gin.Context) { + c.Set("user_id", u.ID) + requestCtx = c.Request.Context() + close(started) + h.List(c) + }) + server := httptest.NewServer(r) + // Also on a regression failure, unlock BEFORE waiting for the server + // to close: the old service otherwise leaves Close waiting on SQL. + t.Cleanup(func() { unlock(); server.Close() }) + clientCtx, cancelClient := context.WithCancel(ctx) + defer cancelClient() + req, err := http.NewRequestWithContext(clientCtx, http.MethodGet, server.URL+"/web/custom-agents", nil) + require.NoError(t, err) + clientDone := make(chan struct{}) + var clientErr error + var status int + requestStart := time.Now() + go func() { + defer close(clientDone) + response, err := server.Client().Do(req) + clientErr = err + if response != nil { + status = response.StatusCode + _, _ = io.Copy(io.Discard, response.Body) + _ = response.Body.Close() + } + }() + testkit.WaitFor(t, 2*time.Second, started, "HTTP handler did not start") + isBlocked := func() bool { + var blocked bool + err := db.WithContext(ctx).Raw("SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = ? AND state = 'active' AND wait_event_type = 'Lock')", queryPID).Scan(&blocked).Error + require.NoError(t, err) + return blocked + } + testkit.Eventually(t, time.Second, isBlocked, "request SQL did not enter the PostgreSQL lock wait", nil) + require.Equal(t, 1, pool.Stats().InUse) + waitsBefore := pool.Stats().WaitCount + pingCtx, cancelPing := context.WithTimeout(ctx, 5*time.Second) + defer cancelPing() + pingDone := make(chan struct{}) + var pingErr error + go func() { defer close(pingDone); pingErr = pool.PingContext(pingCtx) }() + testkit.Eventually(t, time.Second, func() bool { return pool.Stats().WaitCount > waitsBefore }, "second caller did not queue for the single connection", nil) + require.NoError(t, requestCtx.Err(), "fixture must reach pool saturation before cancellation") + + if mode == "disconnect" { + cancelClient() + } + testkit.WaitFor(t, 2*time.Second, requestCtx.Done(), "request context was not canceled") + canceledAt := time.Now() + testkit.Eventually(t, 2*time.Second, func() bool { + select { + case <-serverDone: + default: + return false + } + select { + case <-pingDone: + return true + default: + return false + } + }, "canceled request kept the handler or connection pool occupied while the lock was held", func() string { + return fmt.Sprintf("request_elapsed=%v canceled_elapsed=%v pool=%+v", time.Since(requestStart), time.Since(canceledAt), pool.Stats()) + }) + require.NoError(t, pingErr, "queued caller must acquire a usable connection without unlocking the table") + testkit.Eventually(t, time.Second, func() bool { return !isBlocked() }, "canceled PostgreSQL query is still waiting on the lock", nil) + var lockHeld bool + require.NoError(t, lock.QueryRowContext(ctx, "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE pid = pg_backend_pid() AND relation = 'custom_agents'::regclass AND mode = 'AccessExclusiveLock' AND granted)").Scan(&lockHeld)) + require.True(t, lockHeld, "the fixture must not release its lock to make the pool assertion pass") + assert.Zero(t, pool.Stats().InUse) + testkit.WaitFor(t, time.Second, clientDone, "HTTP client did not finish") + if mode == "deadline" { + assert.ErrorIs(t, requestCtx.Err(), context.DeadlineExceeded) + require.NoError(t, clientErr) + assert.Equal(t, http.StatusGatewayTimeout, status) + } else { + assert.ErrorIs(t, requestCtx.Err(), context.Canceled) + assert.ErrorIs(t, clientErr, context.Canceled) + } + t.Logf("mode=%s deadline=%v request_elapsed=%v canceled_to_released=%v lock_held=%t pool=%+v", mode, deadline, time.Since(requestStart), time.Since(canceledAt), lockHeld, pool.Stats()) + unlock() + agents, err := service.ListCustomAgents(ctx, u.ID) + require.NoError(t, err, "a fresh context must remain usable after cancellation") + require.Len(t, agents, 1) + assert.Equal(t, created.ID, agents[0].ID) + }) + } +}