fix: apply custom prompts to new sessions

Custom prompts saved via the Settings -> Prompts UI are persisted to the
database, but every assistant and flow session creation path was using
templates.NewDefaultPrompter() with a leftover TODO, so user overrides
never reached the agents and Langfuse traces always showed the defaults.

Add a controller-side helper that loads the user's saved prompts and
overlays them on the compiled defaults. Prompt types the user has not
customized continue to use the defaults; an empty body row is treated
as no override (the UI uses delete to reset). A database error fails
session creation explicitly instead of silently falling back.

Wire the helper into the four affected call sites in NewAssistantWorker,
LoadAssistantWorker, NewFlowWorker, and LoadFlowWorker.

Closes #300

Signed-off-by: mason5052 <ehehwnwjs5052@gmail.com>

backend/pkg/controller/assistant.go

@@ -16,7 +16,6 @@ import (
 	"pentagi/pkg/providers"
 	"pentagi/pkg/providers/pconfig"
 	"pentagi/pkg/providers/provider"
-	"pentagi/pkg/templates"
 	"pentagi/pkg/tools"
 
 	"github.com/sirupsen/logrus"
@@ -154,7 +153,10 @@ func NewAssistantWorker(ctx context.Context, awc newAssistantWorkerCtx) (Assista
 		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow assistant log worker", err)
 	}
 
-	prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB
+	prompter, err := newUserPrompter(ctx, awc.db, awc.userID)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to build user prompter", err)
+	}
 	executor, err := tools.NewFlowToolsExecutor(awc.db, awc.cfg, awc.docker, awc.functions, awc.flowID)
 	if err != nil {
 		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow tools executor", err)
@@ -319,7 +321,10 @@ func LoadAssistantWorker(
 		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow assistant log worker", err)
 	}
 
-	prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB
+	prompter, err := newUserPrompter(ctx, awc.db, awc.userID)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to build user prompter", err)
+	}
 	executor, err := tools.NewFlowToolsExecutor(awc.db, awc.cfg, awc.docker, functions, awc.flowID)
 	if err != nil {
 		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow tools executor", err)

backend/pkg/controller/flow.go

@@ -19,7 +19,6 @@ import (
 	"pentagi/pkg/providers"
 	"pentagi/pkg/providers/pconfig"
 	"pentagi/pkg/providers/provider"
-	"pentagi/pkg/templates"
 	"pentagi/pkg/tools"
 
 	"github.com/sirupsen/logrus"
@@ -166,7 +165,10 @@ func NewFlowWorker(
 	flowSpan := observation.Span(langfuse.WithSpanName("prepare flow worker"))
 	ctx, _ = flowSpan.Observation(ctx)
 
-	prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB
+	prompter, err := newUserPrompter(ctx, fwc.db, fwc.userID)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to build user prompter", err)
+	}
 	executor, err := tools.NewFlowToolsExecutor(fwc.db, fwc.cfg, fwc.docker, fwc.functions, flow.ID)
 	if err != nil {
 		return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to create flow tools executor", err)
@@ -332,7 +334,10 @@ func LoadFlowWorker(ctx context.Context, flow database.Flow, fwc flowWorkerCtx)
 		return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to unmarshal functions", err)
 	}
 
-	prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB
+	prompter, err := newUserPrompter(ctx, fwc.db, flow.UserID)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to build user prompter", err)
+	}
 	executor, err := tools.NewFlowToolsExecutor(fwc.db, fwc.cfg, fwc.docker, functions, flow.ID)
 	if err != nil {
 		return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to create flow tools executor", err)

backend/pkg/controller/prompter.go

@@ -0,0 +1,53 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/templates"
+)
+
+// newUserPrompter loads the user's custom prompts from the database and
+// overlays them onto the compiled default templates. Prompt types that
+// the user has not customized continue to use the defaults. A database
+// error is returned to the caller so that session creation fails
+// explicitly instead of silently falling back to defaults.
+func newUserPrompter(ctx context.Context, db database.Querier, userID int64) (templates.Prompter, error) {
+	userPrompts, err := db.GetUserPrompts(ctx, userID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load user prompts: %w", err)
+	}
+	return buildUserPrompter(templates.NewDefaultPrompter(), userPrompts)
+}
+
+// buildUserPrompter is the pure merge step extracted from newUserPrompter
+// so it can be unit-tested without a database fake. It seeds the result
+// with every default template and then overrides each prompt type the
+// user has saved a non-empty body for.
+func buildUserPrompter(defaults templates.Prompter, userPrompts []database.Prompt) (templates.Prompter, error) {
+	blob, err := defaults.DumpTemplates()
+	if err != nil {
+		return nil, fmt.Errorf("failed to dump default templates: %w", err)
+	}
+
+	merged := templates.PromptsMap{}
+	if err := json.Unmarshal(blob, &merged); err != nil {
+		return nil, fmt.Errorf("failed to parse default templates: %w", err)
+	}
+
+	for _, p := range userPrompts {
+		if p.Prompt == "" {
+			// The Prompts UI uses delete (or reset, which writes the
+			// default body back) to remove a customization, so an empty
+			// body is unexpected. Skip it instead of clobbering the
+			// default with an empty string that would later surface as
+			// ErrTemplateNotFound deep inside agent rendering.
+			continue
+		}
+		merged[templates.PromptType(p.Type)] = p.Prompt
+	}
+
+	return templates.NewFlowPrompter(merged), nil
+}

backend/pkg/controller/prompter_test.go

@@ -0,0 +1,211 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/templates"
+)
+
+// fakeQuerier satisfies database.Querier by embedding the interface
+// (so the unused methods stay nil) and overrides only GetUserPrompts.
+// Calling any other method would panic, which is exactly what we want
+// for a unit test that should not touch unrelated DB code.
+type fakeQuerier struct {
+	database.Querier
+	prompts []database.Prompt
+	err     error
+}
+
+func (f *fakeQuerier) GetUserPrompts(ctx context.Context, userID int64) ([]database.Prompt, error) {
+	if f.err != nil {
+		return nil, f.err
+	}
+	return f.prompts, nil
+}
+
+// getDefaultTemplate returns the compiled default body for a prompt
+// type, used to compare overridden vs. preserved prompts in tests.
+func getDefaultTemplate(t *testing.T, pt templates.PromptType) string {
+	t.Helper()
+	body, err := templates.NewDefaultPrompter().GetTemplate(pt)
+	if err != nil {
+		t.Fatalf("default prompter has no template for %q: %v", pt, err)
+	}
+	return body
+}
+
+func TestBuildUserPrompter_NoUserPrompts(t *testing.T) {
+	defaults := templates.NewDefaultPrompter()
+	prompter, err := buildUserPrompter(defaults, nil)
+	if err != nil {
+		t.Fatalf("buildUserPrompter returned error: %v", err)
+	}
+
+	for _, pt := range []templates.PromptType{
+		templates.PromptTypePrimaryAgent,
+		templates.PromptTypeAssistant,
+	} {
+		got, err := prompter.GetTemplate(pt)
+		if err != nil {
+			t.Fatalf("GetTemplate(%q) error: %v", pt, err)
+		}
+		want := getDefaultTemplate(t, pt)
+		if got != want {
+			t.Errorf("prompt %q: expected default body when user has no overrides, got divergent body", pt)
+		}
+	}
+}
+
+func TestBuildUserPrompter_SingleOverride(t *testing.T) {
+	const customBody = "custom primary agent prompt body"
+	userPrompts := []database.Prompt{
+		{Type: database.PromptTypePrimaryAgent, Prompt: customBody},
+	}
+
+	prompter, err := buildUserPrompter(templates.NewDefaultPrompter(), userPrompts)
+	if err != nil {
+		t.Fatalf("buildUserPrompter error: %v", err)
+	}
+
+	got, err := prompter.GetTemplate(templates.PromptTypePrimaryAgent)
+	if err != nil {
+		t.Fatalf("GetTemplate(primary_agent) error: %v", err)
+	}
+	if got != customBody {
+		t.Errorf("primary_agent: expected custom body %q, got %q", customBody, got)
+	}
+
+	// Spot-check that an unrelated type still resolves to the default.
+	for _, pt := range []templates.PromptType{
+		templates.PromptTypeAssistant,
+		templates.PromptTypePentester,
+	} {
+		got, err := prompter.GetTemplate(pt)
+		if err != nil {
+			t.Fatalf("GetTemplate(%q) error: %v", pt, err)
+		}
+		if got != getDefaultTemplate(t, pt) {
+			t.Errorf("prompt %q should still match default after a single unrelated override", pt)
+		}
+	}
+}
+
+func TestBuildUserPrompter_PartialOverridesPreserveDefaults(t *testing.T) {
+	overrides := map[database.PromptType]string{
+		database.PromptTypePrimaryAgent: "custom primary",
+		database.PromptTypeCoder:        "custom coder",
+		database.PromptTypeReporter:     "custom reporter",
+	}
+
+	userPrompts := make([]database.Prompt, 0, len(overrides))
+	for pt, body := range overrides {
+		userPrompts = append(userPrompts, database.Prompt{Type: pt, Prompt: body})
+	}
+
+	prompter, err := buildUserPrompter(templates.NewDefaultPrompter(), userPrompts)
+	if err != nil {
+		t.Fatalf("buildUserPrompter error: %v", err)
+	}
+
+	for pt, want := range overrides {
+		got, err := prompter.GetTemplate(templates.PromptType(pt))
+		if err != nil {
+			t.Fatalf("GetTemplate(%q) error: %v", pt, err)
+		}
+		if got != want {
+			t.Errorf("prompt %q: expected custom body %q, got %q", pt, want, got)
+		}
+	}
+
+	// Types that were not overridden must still match the defaults.
+	for _, pt := range []templates.PromptType{
+		templates.PromptTypeAssistant,
+		templates.PromptTypePentester,
+		templates.PromptTypeSearcher,
+	} {
+		got, err := prompter.GetTemplate(pt)
+		if err != nil {
+			t.Fatalf("GetTemplate(%q) error: %v", pt, err)
+		}
+		if got != getDefaultTemplate(t, pt) {
+			t.Errorf("prompt %q should still match default after partial overrides", pt)
+		}
+	}
+}
+
+func TestBuildUserPrompter_EmptyBodyIsIgnored(t *testing.T) {
+	userPrompts := []database.Prompt{
+		{Type: database.PromptTypePrimaryAgent, Prompt: ""},
+	}
+
+	prompter, err := buildUserPrompter(templates.NewDefaultPrompter(), userPrompts)
+	if err != nil {
+		t.Fatalf("buildUserPrompter error: %v", err)
+	}
+
+	got, err := prompter.GetTemplate(templates.PromptTypePrimaryAgent)
+	if err != nil {
+		t.Fatalf("GetTemplate(primary_agent) error: %v", err)
+	}
+	if got != getDefaultTemplate(t, templates.PromptTypePrimaryAgent) {
+		t.Errorf("primary_agent: empty user body must not clobber default")
+	}
+}
+
+func TestNewUserPrompter_HappyPath(t *testing.T) {
+	const customAssistant = "custom assistant prompt"
+	const customCoder = "custom coder prompt"
+
+	db := &fakeQuerier{
+		prompts: []database.Prompt{
+			{Type: database.PromptTypeAssistant, Prompt: customAssistant},
+			{Type: database.PromptTypeCoder, Prompt: customCoder},
+		},
+	}
+
+	prompter, err := newUserPrompter(context.Background(), db, 42)
+	if err != nil {
+		t.Fatalf("newUserPrompter error: %v", err)
+	}
+
+	for pt, want := range map[templates.PromptType]string{
+		templates.PromptTypeAssistant: customAssistant,
+		templates.PromptTypeCoder:     customCoder,
+	} {
+		got, err := prompter.GetTemplate(pt)
+		if err != nil {
+			t.Fatalf("GetTemplate(%q) error: %v", pt, err)
+		}
+		if got != want {
+			t.Errorf("prompt %q: expected custom body %q, got %q", pt, want, got)
+		}
+	}
+
+	// Sanity check that an un-customized type still falls back.
+	got, err := prompter.GetTemplate(templates.PromptTypePentester)
+	if err != nil {
+		t.Fatalf("GetTemplate(pentester) error: %v", err)
+	}
+	if got != getDefaultTemplate(t, templates.PromptTypePentester) {
+		t.Errorf("pentester: expected default body when user has no override")
+	}
+}
+
+func TestNewUserPrompter_DBErrorPropagates(t *testing.T) {
+	sentinel := errors.New("db connection lost")
+	db := &fakeQuerier{err: sentinel}
+
+	prompter, err := newUserPrompter(context.Background(), db, 42)
+	if err == nil {
+		t.Fatalf("expected error, got nil prompter=%v", prompter)
+	}
+	if prompter != nil {
+		t.Errorf("expected nil prompter on DB error, got %v", prompter)
+	}
+	if !errors.Is(err, sentinel) {
+		t.Errorf("expected wrapped sentinel error, got %v", err)
+	}
+}