fix(controller): sanitize subtask text and persist plans atomically

backend/docs/database.md

@@ -12,6 +12,8 @@ PentAGI stores application state in PostgreSQL and uses the `vector` extension f
 
 SQL queries live in `backend/sqlc/models`. Schema history lives in `backend/migrations/sql`.
 
+Subtask generation, refinement and assistant plan patches use `database.SaveSubtaskPlan` in `pkg/database/subtask_plan.go`. Titles and descriptions pass through `SanitizeUTF8` before insertion: NUL bytes are removed, malformed UTF-8 is replaced, and valid text is preserved. Removal of the replaced subtasks and insertion of the entire new plan share one transaction. Any delete or insert failure rolls back the plan; in-memory subtask workers are invalidated only after commit. A failed or ambiguous commit is reported to the caller. This does not repair partial plans already stored by a previous execution.
+
 This document describes how the product uses the database. It is not a substitute for the schema itself. Sources of truth, in descending order:
 
 1. `backend/migrations/sql/*.sql` — schema and data migrations;

backend/go.mod

@@ -90,6 +90,7 @@ require (
 	cloud.google.com/go/compute/metadata v0.9.0 // indirect
 	cloud.google.com/go/longrunning v0.8.0 // indirect
 	github.com/AssemblyAI/assemblyai-go-sdk v1.3.0 // indirect
+	github.com/DATA-DOG/go-sqlmock v1.5.2
 	github.com/KyleBanks/depth v1.2.1 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/PuerkitoBio/goquery v1.10.3 // indirect

backend/go.sum

@@ -26,6 +26,8 @@ github.com/ClickHouse/ch-go v0.58.2 h1:jSm2szHbT9MCAB1rJ3WuCJqmGLi5UTjlNu+f530UT
 github.com/ClickHouse/ch-go v0.58.2/go.mod h1:Ap/0bEmiLa14gYjCiRkYGbXvbe8vwdrfTYWhsuQ99aw=
 github.com/ClickHouse/clickhouse-go/v2 v2.17.1 h1:ZCmAYWpu75IyEi7+Yrs/uaAjiCGY5wfW5kXo64exkX4=
 github.com/ClickHouse/clickhouse-go/v2 v2.17.1/go.mod h1:rkGTvFDTLqLIm0ma+13xmcCfr/08Gvs7KmFt1tgiWHQ=
+github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU=
+github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU=
 github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc=
 github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE=
 github.com/Masterminds/goutils v1.1.1 h1:5nUrii3FMTL5diU80unEVvNevw1nH4+ZV4DSLVJLSYI=
@@ -373,6 +375,7 @@ github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCV
 github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
+github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE=
 github.com/klauspost/compress v1.10.3/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
 github.com/klauspost/compress v1.18.3 h1:9PJRvfbmTabkOX8moIpXPbMMbYN60bWImDDU7L+/6zw=
 github.com/klauspost/compress v1.18.3/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=

backend/pkg/controller/subtasks.go

@@ -88,19 +88,16 @@ func (stc *subtaskController) GenerateSubtasks(ctx context.Context) error {
 	}
 
 	// TODO: change it to insert subtasks in transaction
+	rows := make([]database.CreateSubtaskParams, 0, len(plan))
 	for _, info := range plan {
-		_, err := stc.taskCtx.DB.CreateSubtask(ctx, database.CreateSubtaskParams{
+		rows = append(rows, database.CreateSubtaskParams{
 			Status:      database.SubtaskStatusCreated,
 			TaskID:      stc.taskCtx.TaskID,
 			Title:       info.Title,
 			Description: info.Description,
 		})
-		if err != nil {
-			return fmt.Errorf("failed to create subtask for task %d: %w", stc.taskCtx.TaskID, err)
-		}
 	}
-
-	return nil
+	return database.SaveSubtaskPlan(ctx, stc.taskCtx.DB, nil, rows)
 }
 
 func (stc *subtaskController) RefineSubtasks(ctx context.Context) error {
@@ -137,24 +134,20 @@ func (stc *subtaskController) RefineSubtasks(ctx context.Context) error {
 		}
 	}
 
-	err = stc.taskCtx.DB.DeleteSubtasks(ctx, subtaskIDs)
-	if err != nil {
-		return fmt.Errorf("failed to delete subtasks for task %d: %w", stc.taskCtx.TaskID, err)
-	}
-	stc.InvalidateSubtasks(subtaskIDs)
-
 	// TODO: change it to insert subtasks in transaction and union it with delete ones
+	rows := make([]database.CreateSubtaskParams, 0, len(plan))
 	for _, info := range plan {
-		_, err := stc.taskCtx.DB.CreateSubtask(ctx, database.CreateSubtaskParams{
+		rows = append(rows, database.CreateSubtaskParams{
 			Status:      database.SubtaskStatusCreated,
 			TaskID:      stc.taskCtx.TaskID,
 			Title:       info.Title,
 			Description: info.Description,
 		})
-		if err != nil {
-			return fmt.Errorf("failed to create subtask for task %d: %w", stc.taskCtx.TaskID, err)
-		}
 	}
+	if err := database.SaveSubtaskPlan(ctx, stc.taskCtx.DB, subtaskIDs, rows); err != nil {
+		return fmt.Errorf("failed to save refined subtasks for task %d: %w", stc.taskCtx.TaskID, err)
+	}
+	stc.InvalidateSubtasks(subtaskIDs)
 
 	return nil
 }

backend/pkg/controller/task_test.go

@@ -83,6 +83,28 @@ func (q *generatorFakeQuerier) GetTaskSubtasks(context.Context, int64) ([]databa
 	return q.subtasks, nil
 }
 
+func (q *generatorFakeQuerier) SaveSubtaskPlan(ctx context.Context, deleteIDs []int64, plan []database.CreateSubtaskParams) error {
+	before := append([]database.Subtask(nil), q.subtasks...)
+	removed := make(map[int64]bool, len(deleteIDs))
+	for _, id := range deleteIDs {
+		removed[id] = true
+	}
+	remaining := make([]database.Subtask, 0, len(q.subtasks))
+	for _, row := range q.subtasks {
+		if !removed[row.ID] {
+			remaining = append(remaining, row)
+		}
+	}
+	q.subtasks = remaining
+	for _, item := range plan {
+		if _, err := q.CreateSubtask(ctx, item); err != nil {
+			q.subtasks = before
+			return err
+		}
+	}
+	return nil
+}
+
 type generatorFakeMsgLog struct {
 	FlowMsgLogWorker
 
@@ -101,10 +123,11 @@ func (m generatorFakeMsgLog) PutTaskMsg(ctx context.Context, _ database.MsglogTy
 type generatorFakeProvider struct {
 	providers.FlowProvider
 
-	plan      []tools.SubtaskInfo
-	err       error
-	cancel    context.CancelFunc
-	refineErr error
+	plan       []tools.SubtaskInfo
+	err        error
+	cancel     context.CancelFunc
+	refineErr  error
+	refinePlan []tools.SubtaskInfo
 }
 
 func (generatorFakeProvider) GetTaskTitle(context.Context, string) (string, error) {
@@ -122,7 +145,7 @@ func (p generatorFakeProvider) GenerateSubtasks(ctx context.Context, _ int64) ([
 }
 
 func (p generatorFakeProvider) RefineSubtasks(context.Context, int64) ([]tools.SubtaskInfo, error) {
-	return nil, p.refineErr
+	return p.refinePlan, p.refineErr
 }
 
 func TestTask_AnExplicitlyEmptyGeneratorPlanRunsTheOriginalRequest(t *testing.T) {
@@ -273,15 +296,16 @@ func TestTask_NewTaskWorker_ClosesATaskThatCannotStart(t *testing.T) {
 		assert.Len(t, pub.taskUpdated, 1)
 	})
 
-	t.Run("a generator that stops part way publishes the subtasks it committed", func(t *testing.T) {
+	t.Run("a rejected generator plan publishes no partial subtasks", func(t *testing.T) {
 		q := &generatorFakeQuerier{subtaskLimit: 2}
 		flowCtx, pub := newContext(q, generatorFakeProvider{plan: []tools.SubtaskInfo{{Title: "a"}, {Title: "b"}, {Title: "c"}}})
 
 		_, err := NewTaskWorker(context.Background(), flowCtx, "input", nil)
 
 		require.ErrorContains(t, err, "subtask row rejected")
 		require.Len(t, pub.taskSubtasks, 1)
-		assert.Len(t, pub.taskSubtasks[0], 2, "the client sees the subtasks the database holds")
+		assert.Empty(t, pub.taskSubtasks[0])
+		assert.Empty(t, q.subtasks)
 	})
 
 	t.Run("a generator stopped by its user closes the task without failing it", func(t *testing.T) {
@@ -299,6 +323,39 @@ func TestTask_NewTaskWorker_ClosesATaskThatCannotStart(t *testing.T) {
 	})
 }
 
+func TestTask_RefinedPlanInvalidatesWorkersOnlyAfterCommit(t *testing.T) {
+	for _, fail := range []bool{false, true} {
+		name := "committed"
+		if fail {
+			name = "rejected"
+		}
+		t.Run(name, func(t *testing.T) {
+			original := []database.Subtask{{ID: 195, Status: database.SubtaskStatusFinished}, {ID: 196, Status: database.SubtaskStatusCreated}}
+			q := &generatorFakeQuerier{subtasks: original}
+			if fail {
+				q.subtaskLimit = 2
+			}
+			stc := NewSubtaskController(&TaskContext{TaskID: 9, FlowContext: FlowContext{DB: q, Provider: generatorFakeProvider{refinePlan: []tools.SubtaskInfo{{Title: "new first"}, {Title: "new second"}}}}}).(*subtaskController)
+			stc.subtasks[195] = nil
+			stc.subtasks[196] = nil
+			err := stc.RefineSubtasks(context.Background())
+			_, oldWorkerExists := stc.subtasks[196]
+			_, completedWorkerExists := stc.subtasks[195]
+			require.True(t, completedWorkerExists)
+			if fail {
+				require.Error(t, err)
+				assert.Equal(t, original, q.subtasks)
+				assert.True(t, oldWorkerExists)
+			} else {
+				require.NoError(t, err)
+				assert.False(t, oldWorkerExists)
+				require.Len(t, q.subtasks, 3)
+				assert.Equal(t, original[0], q.subtasks[0])
+			}
+		})
+	}
+}
+
 func TestTask_Report_RefusesACompletedTaskWithASentinel(t *testing.T) {
 	t.Parallel()
 

backend/pkg/database/subtask_plan.go

@@ -0,0 +1,51 @@
+package database
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+)
+
+// SaveSubtaskPlan commits deletion and creation together. Text is normalized
+// at this boundary because PostgreSQL text cannot contain NUL bytes.
+func SaveSubtaskPlan(ctx context.Context, db Querier, deleteIDs []int64, plan []CreateSubtaskParams) error {
+	writer, ok := db.(interface {
+		SaveSubtaskPlan(context.Context, []int64, []CreateSubtaskParams) error
+	})
+	if !ok {
+		return errors.New("subtask plan storage does not support transactions")
+	}
+	return writer.SaveSubtaskPlan(ctx, deleteIDs, plan)
+}
+
+func (q *Queries) SaveSubtaskPlan(ctx context.Context, deleteIDs []int64, plan []CreateSubtaskParams) error {
+	beginner, ok := q.db.(interface {
+		BeginTx(context.Context, *sql.TxOptions) (*sql.Tx, error)
+	})
+	if !ok {
+		return errors.New("subtask plan storage cannot begin a transaction")
+	}
+	tx, err := beginner.BeginTx(ctx, nil)
+	if err != nil {
+		return fmt.Errorf("begin subtask plan transaction: %w", err)
+	}
+	defer func() { _ = tx.Rollback() }()
+	queries := q.WithTx(tx)
+	if len(deleteIDs) > 0 {
+		if err := queries.DeleteSubtasks(ctx, deleteIDs); err != nil {
+			return fmt.Errorf("delete planned subtasks: %w", err)
+		}
+	}
+	for _, item := range plan {
+		item.Title = SanitizeUTF8(item.Title)
+		item.Description = SanitizeUTF8(item.Description)
+		if _, err := queries.CreateSubtask(ctx, item); err != nil {
+			return fmt.Errorf("create subtask for task %d: %w", item.TaskID, err)
+		}
+	}
+	if err := tx.Commit(); err != nil {
+		return fmt.Errorf("commit subtask plan: %w", err)
+	}
+	return nil
+}

backend/pkg/database/subtask_plan_test.go

@@ -0,0 +1,84 @@
+package database
+
+import (
+	"context"
+	"errors"
+	"regexp"
+	"testing"
+	"time"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/lib/pq"
+	"github.com/stretchr/testify/require"
+)
+
+func subtaskPlanRow(id int64, item CreateSubtaskParams) *sqlmock.Rows {
+	return sqlmock.NewRows([]string{"id", "status", "title", "description", "result", "task_id", "created_at", "updated_at", "context"}).
+		AddRow(id, item.Status, item.Title, item.Description, "", item.TaskID, time.Now(), time.Now(), "")
+}
+
+func TestSaveSubtaskPlanNormalizesTextAndCommits(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	defer db.Close()
+	item := CreateSubtaskParams{TaskID: 9, Status: SubtaskStatusCreated, Title: "验\x00证", Description: "证据\x00\xff\n保留"}
+	cleaned := item
+	cleaned.Title, cleaned.Description = "验证", "证据�\n保留"
+	mock.ExpectBegin()
+	mock.ExpectExec(regexp.QuoteMeta(deleteSubtasks)).WithArgs(pq.Array([]int64{196, 197})).WillReturnResult(sqlmock.NewResult(0, 2))
+	mock.ExpectQuery(regexp.QuoteMeta(createSubtask)).WithArgs(item.Status, cleaned.Title, cleaned.Description, item.TaskID).WillReturnRows(subtaskPlanRow(201, cleaned))
+	mock.ExpectCommit()
+	require.NoError(t, SaveSubtaskPlan(context.Background(), New(db), []int64{196, 197}, []CreateSubtaskParams{item}))
+	require.Equal(t, "验\x00证", item.Title)
+	require.NoError(t, mock.ExpectationsWereMet())
+}
+
+func TestSaveSubtaskPlanRollsBackPartialWrites(t *testing.T) {
+	for _, replace := range []bool{false, true} {
+		name := "generation"
+		if replace {
+			name = "refinement"
+		}
+		t.Run(name, func(t *testing.T) {
+			db, mock, err := sqlmock.New()
+			require.NoError(t, err)
+			defer db.Close()
+			item := CreateSubtaskParams{TaskID: 9, Status: SubtaskStatusCreated, Title: "first", Description: "description"}
+			failure := errors.New("second row rejected")
+			mock.ExpectBegin()
+			var ids []int64
+			if replace {
+				ids = []int64{196, 197}
+				mock.ExpectExec(regexp.QuoteMeta(deleteSubtasks)).WithArgs(pq.Array(ids)).WillReturnResult(sqlmock.NewResult(0, 2))
+			}
+			mock.ExpectQuery(regexp.QuoteMeta(createSubtask)).WithArgs(item.Status, item.Title, item.Description, item.TaskID).WillReturnRows(subtaskPlanRow(201, item))
+			mock.ExpectQuery(regexp.QuoteMeta(createSubtask)).WillReturnError(failure)
+			mock.ExpectRollback()
+			require.ErrorIs(t, SaveSubtaskPlan(context.Background(), New(db), ids, []CreateSubtaskParams{item, item}), failure)
+			require.NoError(t, mock.ExpectationsWereMet())
+		})
+	}
+}
+
+func TestSaveSubtaskPlanDeleteFailureRollsBack(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	defer db.Close()
+	failure := errors.New("delete rejected")
+	mock.ExpectBegin()
+	mock.ExpectExec(regexp.QuoteMeta(deleteSubtasks)).WillReturnError(failure)
+	mock.ExpectRollback()
+	require.ErrorIs(t, SaveSubtaskPlan(context.Background(), New(db), []int64{196}, nil), failure)
+	require.NoError(t, mock.ExpectationsWereMet())
+}
+
+func TestSaveSubtaskPlanCommitFailureIsReturned(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	defer db.Close()
+	failure := errors.New("commit failed")
+	mock.ExpectBegin()
+	mock.ExpectCommit().WillReturnError(failure)
+	require.ErrorIs(t, SaveSubtaskPlan(context.Background(), New(db), nil, nil), failure)
+	require.NoError(t, mock.ExpectationsWereMet())
+}

backend/pkg/providers/assistant.go

@@ -647,27 +647,22 @@ func (ap *assistantProvider) patchAssistantFlowSubtasks(
 		idsToDelete = append(idsToDelete, st.ID)
 	}
 
-	if len(idsToDelete) > 0 {
-		if err := db.DeleteSubtasks(ctx, idsToDelete); err != nil {
-			return fmt.Errorf("failed to delete subtasks for task %d: %w", taskID, err)
-		}
-		// The task worker for taskID (if currently loaded in memory) still holds
-		// SubtaskWorker objects for idsToDelete; without this it would later try
-		// to operate on rows that no longer exist (e.g. on flow Finish/Stop).
-		ap.flowWorker.InvalidateTaskSubtasks(ctx, taskID, idsToDelete)
-	}
-
 	// result contains only entries that were in patchableSubs (minus removed ones)
 	// plus any new entries from add operations — all safe to create unconditionally.
+	rows := make([]database.CreateSubtaskParams, 0, len(result))
 	for _, info := range result {
-		if _, err := db.CreateSubtask(ctx, database.CreateSubtaskParams{
+		rows = append(rows, database.CreateSubtaskParams{
 			Status:      database.SubtaskStatusCreated,
 			TaskID:      taskID,
 			Title:       info.Title,
 			Description: info.Description,
-		}); err != nil {
-			return fmt.Errorf("failed to create subtask %q for task %d: %w", info.Title, taskID, err)
-		}
+		})
+	}
+	if err := database.SaveSubtaskPlan(ctx, db, idsToDelete, rows); err != nil {
+		return fmt.Errorf("failed to save patched subtasks for task %d: %w", taskID, err)
+	}
+	if len(idsToDelete) > 0 {
+		ap.flowWorker.InvalidateTaskSubtasks(ctx, taskID, idsToDelete)
 	}
 
 	subtasksResult := convertSubtaskInfoPatch(result)