chore: import upstream snapshot with attribution

backend/pkg/controller/alog.go

@@ -0,0 +1,81 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type FlowAgentLogWorker interface {
+	PutLog(
+		ctx context.Context,
+		initiator database.MsgchainType,
+		executor database.MsgchainType,
+		task string,
+		result string,
+		taskID *int64,
+		subtaskID *int64,
+	) (int64, error)
+	GetLog(ctx context.Context, msgID int64) (database.Agentlog, error)
+}
+
+type flowAgentLogWorker struct {
+	db     database.Querier
+	mx     *sync.Mutex
+	flowID int64
+	pub    subscriptions.FlowPublisher
+}
+
+func NewFlowAgentLogWorker(db database.Querier, flowID int64, pub subscriptions.FlowPublisher) FlowAgentLogWorker {
+	return &flowAgentLogWorker{
+		db:     db,
+		mx:     &sync.Mutex{},
+		flowID: flowID,
+		pub:    pub,
+	}
+}
+
+func (flw *flowAgentLogWorker) PutLog(
+	ctx context.Context,
+	initiator database.MsgchainType,
+	executor database.MsgchainType,
+	task string,
+	result string,
+	taskID *int64,
+	subtaskID *int64,
+) (int64, error) {
+	flw.mx.Lock()
+	defer flw.mx.Unlock()
+
+	flLog, err := flw.db.CreateAgentLog(ctx, database.CreateAgentLogParams{
+		Initiator: initiator,
+		Executor:  executor,
+		Task:      task,
+		Result:    result,
+		FlowID:    flw.flowID,
+		TaskID:    database.Int64ToNullInt64(taskID),
+		SubtaskID: database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to create search log: %w", err)
+	}
+
+	flw.pub.AgentLogAdded(ctx, flLog)
+
+	return flLog.ID, nil
+}
+
+func (flw *flowAgentLogWorker) GetLog(ctx context.Context, msgID int64) (database.Agentlog, error) {
+	msg, err := flw.db.GetFlowAgentLog(ctx, database.GetFlowAgentLogParams{
+		ID:     msgID,
+		FlowID: flw.flowID,
+	})
+	if err != nil {
+		return database.Agentlog{}, fmt.Errorf("failed to get agent log: %w", err)
+	}
+
+	return msg, nil
+}

backend/pkg/controller/alogs.go

@@ -0,0 +1,71 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type AgentLogController interface {
+	NewFlowAgentLog(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowAgentLogWorker, error)
+	ListFlowsAgentLog(ctx context.Context) ([]FlowAgentLogWorker, error)
+	GetFlowAgentLog(ctx context.Context, flowID int64) (FlowAgentLogWorker, error)
+}
+
+type agentLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowAgentLogWorker
+}
+
+func NewAgentLogController(db database.Querier) AgentLogController {
+	return &agentLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowAgentLogWorker),
+	}
+}
+
+func (alc *agentLogController) NewFlowAgentLog(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowAgentLogWorker, error) {
+	alc.mx.Lock()
+	defer alc.mx.Unlock()
+
+	flw := NewFlowAgentLogWorker(alc.db, flowID, pub)
+	alc.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (alc *agentLogController) ListFlowsAgentLog(ctx context.Context) ([]FlowAgentLogWorker, error) {
+	alc.mx.Lock()
+	defer alc.mx.Unlock()
+
+	flows := make([]FlowAgentLogWorker, 0, len(alc.flows))
+	for _, flw := range alc.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (alc *agentLogController) GetFlowAgentLog(
+	ctx context.Context,
+	flowID int64,
+) (FlowAgentLogWorker, error) {
+	alc.mx.Lock()
+	defer alc.mx.Unlock()
+
+	flw, ok := alc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}

backend/pkg/controller/aslog.go

@@ -0,0 +1,428 @@
+package controller
+
+import (
+	"bytes"
+	"context"
+	"sync"
+	"time"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+	"pentagi/pkg/providers"
+
+	lru "github.com/hashicorp/golang-lru/v2/expirable"
+	"github.com/vxcontrol/langchaingo/llms/reasoning"
+)
+
+const (
+	updateMsgTimeout = 30 * time.Second
+	streamCacheSize  = 1000
+	streamCacheTTL   = 2 * time.Hour
+)
+
+type FlowAssistantLogWorker interface {
+	PutMsg(
+		ctx context.Context,
+		msgType database.MsglogType,
+		taskID, subtaskID *int64,
+		streamID int64,
+		thinking, msg string,
+	) (int64, error)
+	PutFlowAssistantMsg(
+		ctx context.Context,
+		msgType database.MsglogType,
+		thinking, msg string,
+	) (int64, error)
+	PutFlowAssistantMsgResult(
+		ctx context.Context,
+		msgType database.MsglogType,
+		thinking, msg, result string,
+		resultFormat database.MsglogResultFormat,
+	) (int64, error)
+	StreamFlowAssistantMsg(
+		ctx context.Context,
+		chunk *providers.StreamMessageChunk,
+	) error
+	UpdateMsgResult(
+		ctx context.Context,
+		msgID, streamID int64,
+		result string,
+		resultFormat database.MsglogResultFormat,
+	) error
+}
+
+type flowAssistantLogWorker struct {
+	db          database.Querier
+	mx          *sync.Mutex
+	flowID      int64
+	assistantID int64
+	results     map[int64]chan *providers.StreamMessageChunk
+	streamCache *lru.LRU[int64, int64] // streamID -> msgID
+	pub         subscriptions.FlowPublisher
+}
+
+func NewFlowAssistantLogWorker(
+	db database.Querier, flowID int64, assistantID int64, pub subscriptions.FlowPublisher,
+) FlowAssistantLogWorker {
+	return &flowAssistantLogWorker{
+		db:          db,
+		mx:          &sync.Mutex{},
+		flowID:      flowID,
+		assistantID: assistantID,
+		results:     make(map[int64]chan *providers.StreamMessageChunk),
+		streamCache: lru.NewLRU[int64, int64](streamCacheSize, nil, streamCacheTTL),
+		pub:         pub,
+	}
+}
+
+func (aslw *flowAssistantLogWorker) PutMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID *int64,
+	streamID int64,
+	thinking, msg string,
+) (int64, error) {
+	aslw.mx.Lock()
+	defer aslw.mx.Unlock()
+
+	return aslw.putMsg(ctx, msgType, taskID, subtaskID, streamID, thinking, msg)
+}
+
+func (aslw *flowAssistantLogWorker) PutFlowAssistantMsg(
+	ctx context.Context, msgType database.MsglogType, thinking, msg string,
+) (int64, error) {
+	aslw.mx.Lock()
+	defer aslw.mx.Unlock()
+
+	return aslw.putMsg(ctx, msgType, nil, nil, 0, thinking, msg)
+}
+
+func (aslw *flowAssistantLogWorker) PutFlowAssistantMsgResult(
+	ctx context.Context, msgType database.MsglogType, thinking, msg, result string,
+	resultFormat database.MsglogResultFormat,
+) (int64, error) {
+	aslw.mx.Lock()
+	defer aslw.mx.Unlock()
+
+	return aslw.putMsgResult(ctx, msgType, nil, nil, thinking, msg, result, resultFormat)
+}
+
+func (aslw *flowAssistantLogWorker) StreamFlowAssistantMsg(
+	ctx context.Context, chunk *providers.StreamMessageChunk,
+) error {
+	aslw.mx.Lock()
+	defer aslw.mx.Unlock()
+
+	return aslw.appendMsgResult(ctx, chunk)
+}
+
+func (aslw *flowAssistantLogWorker) UpdateMsgResult(
+	ctx context.Context,
+	msgID, streamID int64,
+	result string,
+	resultFormat database.MsglogResultFormat,
+) error {
+	aslw.mx.Lock()
+	defer aslw.mx.Unlock()
+
+	msgLog, err := aslw.db.GetFlowAssistantLog(ctx, msgID)
+	if err != nil {
+		return err
+	}
+
+	ch, workerFound := aslw.results[streamID]
+	if workerFound {
+		ch <- &providers.StreamMessageChunk{
+			Type:         providers.StreamMessageChunkTypeResult,
+			MsgType:      msgLog.Type,
+			Content:      msgLog.Message,
+			Thinking:     aslw.getThinkingStructure(msgLog.Thinking.String),
+			Result:       result,
+			ResultFormat: resultFormat,
+			StreamID:     streamID,
+		}
+		return nil
+	}
+
+	msgLog, err = aslw.db.UpdateAssistantLogResult(ctx, database.UpdateAssistantLogResultParams{
+		Result:       database.SanitizeUTF8(result),
+		ResultFormat: resultFormat,
+		ID:           msgID,
+	})
+	if err != nil {
+		return err
+	}
+
+	aslw.pub.AssistantLogUpdated(ctx, msgLog, false)
+
+	return nil
+}
+
+func (aslw *flowAssistantLogWorker) putMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID *int64,
+	streamID int64,
+	thinking, msg string,
+) (int64, error) {
+	msgID, msgFound := aslw.streamCache.Get(streamID)
+	ch, workerFound := aslw.results[streamID]
+
+	if msgFound && workerFound {
+		ch <- &providers.StreamMessageChunk{
+			Type:     providers.StreamMessageChunkTypeUpdate,
+			MsgType:  msgType,
+			Content:  msg,
+			Thinking: aslw.getThinkingStructure(thinking),
+			StreamID: streamID,
+		}
+		return msgID, nil
+	} else if msgFound {
+		msgLog, err := aslw.db.UpdateAssistantLogContent(ctx, database.UpdateAssistantLogContentParams{
+			Type:     msgType,
+			Message:  database.SanitizeUTF8(msg),
+			Thinking: database.StringToNullString(database.SanitizeUTF8(thinking)),
+			ID:       msgID,
+		})
+		if err == nil {
+			aslw.pub.AssistantLogUpdated(ctx, msgLog, false)
+		}
+		return msgID, err
+	} else {
+		msgLog, err := aslw.db.CreateAssistantLog(ctx, database.CreateAssistantLogParams{
+			Type:        msgType,
+			Message:     database.SanitizeUTF8(msg),
+			Thinking:    database.StringToNullString(database.SanitizeUTF8(thinking)),
+			FlowID:      aslw.flowID,
+			AssistantID: aslw.assistantID,
+		})
+		if err == nil {
+			if streamID != 0 {
+				aslw.streamCache.Add(streamID, msgLog.ID)
+			}
+			aslw.pub.AssistantLogAdded(ctx, msgLog)
+			return msgLog.ID, nil
+		}
+		return 0, err
+	}
+}
+
+func (aslw *flowAssistantLogWorker) putMsgResult(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID *int64,
+	thinking, msg, result string,
+	resultFormat database.MsglogResultFormat,
+) (int64, error) {
+	if len(msg) > defaultMaxMessageLength {
+		msg = msg[:defaultMaxMessageLength] + "..."
+	}
+
+	msgLog, err := aslw.db.CreateResultAssistantLog(ctx, database.CreateResultAssistantLogParams{
+		Type:         msgType,
+		Message:      database.SanitizeUTF8(msg),
+		Thinking:     database.StringToNullString(database.SanitizeUTF8(thinking)),
+		Result:       database.SanitizeUTF8(result),
+		ResultFormat: resultFormat,
+		FlowID:       aslw.flowID,
+		AssistantID:  aslw.assistantID,
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	aslw.pub.AssistantLogAdded(ctx, msgLog)
+
+	return msgLog.ID, nil
+}
+
+func (aslw *flowAssistantLogWorker) appendMsgResult(
+	ctx context.Context, chunk *providers.StreamMessageChunk,
+) error {
+	var (
+		err    error
+		msgLog database.Assistantlog
+	)
+
+	if chunk == nil {
+		return nil
+	}
+
+	ch, ok := aslw.results[chunk.StreamID]
+	if ok {
+		ch <- chunk
+		return nil
+	}
+
+	msgLog, err = aslw.db.CreateAssistantLog(ctx, database.CreateAssistantLogParams{
+		Type:        chunk.MsgType,
+		Message:     "", // special case for completion answer
+		FlowID:      aslw.flowID,
+		AssistantID: aslw.assistantID,
+	})
+	if err != nil {
+		return err
+	}
+
+	aslw.streamCache.Add(chunk.StreamID, msgLog.ID)
+	ch = make(chan *providers.StreamMessageChunk, 50) // safe capacity to avoid deadlock
+	aslw.results[chunk.StreamID] = ch                 // it's safe because mutex is used in parent method
+	ch <- chunk
+
+	go aslw.workerMsgUpdater(msgLog.ID, chunk.StreamID, ch)
+
+	return nil
+}
+
+func (aslw *flowAssistantLogWorker) workerMsgUpdater(
+	msgID, streamID int64,
+	ch chan *providers.StreamMessageChunk,
+) {
+	timer := time.NewTimer(updateMsgTimeout)
+	defer timer.Stop()
+
+	ctx := context.Background()
+	result := ""
+	resultFormat := database.MsglogResultFormatPlain
+	contentData := make([]byte, 0, defaultMaxMessageLength)
+	contentBuf := bytes.NewBuffer(contentData)
+	thinkingData := make([]byte, 0, defaultMaxMessageLength)
+	thinkingBuf := bytes.NewBuffer(thinkingData)
+	wasUpdated := false // track if we actually updated the record
+
+	msgLog, err := aslw.db.GetFlowAssistantLog(ctx, msgID)
+	if err != nil {
+		// generic fields
+		msgLog = database.Assistantlog{
+			ID:          msgID,
+			FlowID:      aslw.flowID,
+			AssistantID: aslw.assistantID,
+			CreatedAt:   database.TimeToNullTime(time.Now()),
+		}
+	}
+
+	newLog := func(msgType database.MsglogType, content, thinking string) database.Assistantlog {
+		return database.Assistantlog{
+			ID:           msgID,
+			Type:         msgType,
+			Message:      content,
+			Thinking:     database.StringToNullString(thinking),
+			Result:       result,
+			ResultFormat: resultFormat,
+			FlowID:       msgLog.FlowID,
+			AssistantID:  msgLog.AssistantID,
+			CreatedAt:    msgLog.CreatedAt,
+		}
+	}
+
+	processChunk := func(chunk *providers.StreamMessageChunk) {
+		switch chunk.Type {
+		case providers.StreamMessageChunkTypeUpdate:
+			thinkingBuf.Reset()
+			contentBuf.Reset()
+			thinkingBuf.WriteString(aslw.getThinkingString(chunk.Thinking))
+			contentBuf.WriteString(chunk.Content)
+			fallthrough // update both thinking and content, send it via publisher
+
+		case providers.StreamMessageChunkTypeFlush:
+			content, thinking := contentBuf.String(), thinkingBuf.String()
+			msgLog, err = aslw.db.UpdateAssistantLogContent(ctx, database.UpdateAssistantLogContentParams{
+				Type:     chunk.MsgType,
+				Message:  database.SanitizeUTF8(content),
+				Thinking: database.StringToNullString(database.SanitizeUTF8(thinking)),
+				ID:       msgID,
+			})
+			if err == nil {
+				wasUpdated = true
+				aslw.pub.AssistantLogUpdated(ctx, msgLog, false)
+			}
+
+		case providers.StreamMessageChunkTypeContent:
+			contentBuf.WriteString(chunk.Content)
+			wasUpdated = true
+			aslw.pub.AssistantLogUpdated(ctx, newLog(chunk.MsgType, chunk.Content, ""), true)
+
+		case providers.StreamMessageChunkTypeThinking:
+			thinkingBuf.WriteString(aslw.getThinkingString(chunk.Thinking))
+			wasUpdated = true
+			aslw.pub.AssistantLogUpdated(ctx, newLog(chunk.MsgType, "", aslw.getThinkingString(chunk.Thinking)), true)
+
+		case providers.StreamMessageChunkTypeResult:
+			result = chunk.Result
+			resultFormat = chunk.ResultFormat
+			content, thinking := contentBuf.String(), thinkingBuf.String()
+			msgLog, err = aslw.db.UpdateAssistantLog(ctx, database.UpdateAssistantLogParams{
+				Type:         chunk.MsgType,
+				Message:      database.SanitizeUTF8(content),
+				Thinking:     database.StringToNullString(database.SanitizeUTF8(thinking)),
+				Result:       database.SanitizeUTF8(result),
+				ResultFormat: resultFormat,
+				ID:           msgID,
+			})
+			if err == nil {
+				wasUpdated = true
+				aslw.pub.AssistantLogUpdated(ctx, msgLog, false)
+			}
+		}
+	}
+
+	for {
+		select {
+		case <-timer.C:
+			// TryLock avoids a deadlock: while this branch holds up reading from ch,
+			// StreamFlowAssistantMsg may be blocked on a full ch <- chunk while
+			// holding aslw.mx. If the mutex is taken, the stream is still active —
+			// reset the timer and keep draining instead of blocking.
+			if !aslw.mx.TryLock() {
+				timer.Reset(updateMsgTimeout)
+				continue
+			}
+			defer aslw.mx.Unlock()
+
+			for i := 0; i < len(ch); i++ {
+				processChunk(<-ch)
+			}
+
+			// If record was never updated, delete it (empty message case)
+			if !wasUpdated {
+				_ = aslw.db.DeleteFlowAssistantLog(ctx, msgID)
+			} else if msgLog, err = aslw.db.GetFlowAssistantLog(ctx, msgID); err == nil {
+				content, thinking := contentBuf.String(), thinkingBuf.String()
+				_, _ = aslw.db.UpdateAssistantLog(ctx, database.UpdateAssistantLogParams{
+					Type:         msgLog.Type,
+					Message:      database.SanitizeUTF8(content),
+					Thinking:     database.StringToNullString(database.SanitizeUTF8(thinking)),
+					Result:       msgLog.Result,
+					ResultFormat: msgLog.ResultFormat,
+					ID:           msgID,
+				})
+			}
+
+			delete(aslw.results, streamID)
+			close(ch)
+
+			return
+
+		case chunk := <-ch:
+			timer.Reset(updateMsgTimeout)
+
+			processChunk(chunk)
+		}
+	}
+}
+
+func (aslw *flowAssistantLogWorker) getThinkingString(thinking *reasoning.ContentReasoning) string {
+	if thinking == nil {
+		return ""
+	}
+	return thinking.Content
+}
+
+func (aslw *flowAssistantLogWorker) getThinkingStructure(thinking string) *reasoning.ContentReasoning {
+	if thinking == "" {
+		return nil
+	}
+	return &reasoning.ContentReasoning{
+		Content: thinking,
+	}
+}

backend/pkg/controller/aslogs.go

@@ -0,0 +1,84 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type AssistantLogController interface {
+	NewFlowAssistantLog(
+		ctx context.Context, flowID int64, assistantID int64, pub subscriptions.FlowPublisher,
+	) (FlowAssistantLogWorker, error)
+	ListFlowsAssistantLog(ctx context.Context, flowID int64) ([]FlowAssistantLogWorker, error)
+	GetFlowAssistantLog(ctx context.Context, flowID int64, assistantID int64) (FlowAssistantLogWorker, error)
+}
+
+type assistantLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]map[int64]FlowAssistantLogWorker
+}
+
+func NewAssistantLogController(db database.Querier) AssistantLogController {
+	return &assistantLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]map[int64]FlowAssistantLogWorker),
+	}
+}
+
+func (aslc *assistantLogController) NewFlowAssistantLog(
+	ctx context.Context, flowID, assistantID int64, pub subscriptions.FlowPublisher,
+) (FlowAssistantLogWorker, error) {
+	aslc.mx.Lock()
+	defer aslc.mx.Unlock()
+
+	flw := NewFlowAssistantLogWorker(aslc.db, flowID, assistantID, pub)
+	if _, ok := aslc.flows[flowID]; !ok {
+		aslc.flows[flowID] = make(map[int64]FlowAssistantLogWorker)
+	}
+	aslc.flows[flowID][assistantID] = flw
+
+	return flw, nil
+}
+
+func (aslc *assistantLogController) ListFlowsAssistantLog(
+	ctx context.Context, flowID int64,
+) ([]FlowAssistantLogWorker, error) {
+	aslc.mx.Lock()
+	defer aslc.mx.Unlock()
+
+	if _, ok := aslc.flows[flowID]; !ok {
+		return []FlowAssistantLogWorker{}, nil
+	}
+
+	flows := make([]FlowAssistantLogWorker, 0, len(aslc.flows[flowID]))
+	for _, flw := range aslc.flows[flowID] {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (aslc *assistantLogController) GetFlowAssistantLog(
+	ctx context.Context, flowID, assistantID int64,
+) (FlowAssistantLogWorker, error) {
+	aslc.mx.Lock()
+	defer aslc.mx.Unlock()
+
+	flw, ok := aslc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	aslw, ok := flw[assistantID]
+	if !ok {
+		return nil, fmt.Errorf("assistant not found")
+	}
+
+	return aslw, nil
+}

backend/pkg/controller/assistant.go

@@ -0,0 +1,614 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"pentagi/pkg/cast"
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+	obs "pentagi/pkg/observability"
+	"pentagi/pkg/observability/langfuse"
+	"pentagi/pkg/providers"
+	"pentagi/pkg/providers/pconfig"
+	"pentagi/pkg/providers/provider"
+	"pentagi/pkg/templates"
+	"pentagi/pkg/tools"
+
+	"github.com/sirupsen/logrus"
+)
+
+const stopAssistantTimeout = 5 * time.Second
+
+type AssistantWorker interface {
+	GetAssistantID() int64
+	GetUserID() int64
+	GetFlowID() int64
+	GetTitle() string
+	GetStatus(ctx context.Context) (database.AssistantStatus, error)
+	SetStatus(ctx context.Context, status database.AssistantStatus) error
+	PutInput(ctx context.Context, input string, useAgents bool, resources []database.UserResource) error
+	Finish(ctx context.Context) error
+	Stop(ctx context.Context) error
+}
+
+type assistantWorker struct {
+	id      int64
+	flowID  int64
+	userID  int64
+	chainID int64
+	aslw    FlowAssistantLogWorker
+	ap      providers.AssistantProvider
+	db      database.Querier
+	wg      *sync.WaitGroup
+	pub     subscriptions.FlowPublisher
+	fw      FlowWorker
+	ctx     context.Context
+	cancel  context.CancelFunc
+	runMX   *sync.Mutex
+	runST   context.CancelFunc
+	runWG   *sync.WaitGroup
+	input   chan assistantInput
+	logger  *logrus.Entry
+}
+
+type newAssistantWorkerCtx struct {
+	userID    int64
+	flowID    int64
+	input     string
+	useAgents bool
+	prvname   provider.ProviderName
+	prvtype   provider.ProviderType
+	functions *tools.Functions
+	resources []database.UserResource
+	fw        FlowWorker
+
+	flowWorkerCtx
+}
+
+type assistantWorkerCtx struct {
+	userID int64
+	flowID int64
+	fw     FlowWorker
+
+	// prompter is an optional reusable prompter for the user. When non-nil,
+	// LoadAssistantWorker uses it instead of issuing another GetUserPrompts
+	// query and merging defaults. LoadFlowWorker sets this so a flow load
+	// with multiple assistants only pays the DB+merge cost once.
+	prompter templates.Prompter
+
+	flowWorkerCtx
+}
+
+const assistantInputTimeout = 2 * time.Second
+
+type assistantInput struct {
+	input     string
+	useAgents bool
+	done      chan error
+}
+
+func NewAssistantWorker(ctx context.Context, awc newAssistantWorkerCtx) (AssistantWorker, error) {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.NewAssistantWorker")
+	defer span.End()
+
+	logger := logrus.WithContext(ctx).WithFields(logrus.Fields{
+		"flow_id":       awc.flowID,
+		"user_id":       awc.userID,
+		"provider_name": awc.prvname.String(),
+		"provider_type": awc.prvtype.String(),
+	})
+
+	user, err := awc.db.GetUser(ctx, awc.userID)
+	if err != nil {
+		logger.WithError(err).Error("failed to get user")
+		return nil, fmt.Errorf("failed to get user %d: %w", awc.userID, err)
+	}
+
+	container, err := awc.db.GetFlowPrimaryContainer(ctx, awc.flowID)
+	if err != nil {
+		logger.WithError(err).Error("failed to get flow primary container")
+		return nil, fmt.Errorf("failed to get flow primary container: %w", err)
+	}
+
+	assistant, err := awc.db.CreateAssistant(ctx, database.CreateAssistantParams{
+		Title:              "untitled",
+		Status:             database.AssistantStatusCreated,
+		Model:              "unknown",
+		ModelProviderName:  string(awc.prvname),
+		ModelProviderType:  database.ProviderType(awc.prvtype),
+		Language:           "English",
+		ToolCallIDTemplate: cast.ToolCallIDTemplate,
+		Functions:          []byte("{}"),
+		FlowID:             awc.flowID,
+		UseAgents:          awc.useAgents,
+	})
+	if err != nil {
+		logger.WithError(err).Error("failed to create assistant in DB")
+		return nil, fmt.Errorf("failed to create assistant in DB: %w", err)
+	}
+
+	logger = logger.WithField("assistant_id", assistant.ID)
+	logger.Info("assistant created in DB")
+
+	ctx, observation := obs.Observer.NewObservation(ctx,
+		langfuse.WithObservationTraceContext(
+			langfuse.WithTraceName(fmt.Sprintf("%d flow %d assistant worker", awc.flowID, assistant.ID)),
+			langfuse.WithTraceUserID(user.Mail),
+			langfuse.WithTraceTags([]string{"controller", "assistant"}),
+			langfuse.WithTraceInput(awc.input),
+			langfuse.WithTraceSessionID(fmt.Sprintf("assistant-%d-flow-%d", assistant.ID, awc.flowID)),
+			langfuse.WithTraceMetadata(langfuse.Metadata{
+				"assistant_id":  assistant.ID,
+				"flow_id":       awc.flowID,
+				"user_id":       awc.userID,
+				"user_email":    user.Mail,
+				"user_name":     user.Name,
+				"user_hash":     user.Hash,
+				"user_role":     user.RoleName,
+				"provider_name": awc.prvname.String(),
+				"provider_type": awc.prvtype.String(),
+			}),
+		),
+	)
+	assistantSpan := observation.Span(langfuse.WithSpanName("prepare assistant worker"))
+	ctx, _ = assistantSpan.Observation(ctx)
+
+	pub := awc.subs.NewFlowPublisher(awc.userID, awc.flowID)
+	aslw, err := awc.aslc.NewFlowAssistantLog(ctx, awc.flowID, assistant.ID, pub)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow assistant log worker", err)
+	}
+
+	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.userID, awc.flowID)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow tools executor", err)
+	}
+	assistantProvider, err := awc.provs.NewAssistantProvider(ctx, awc.prvname, prompter, executor,
+		assistant.ID, awc.flowID, awc.userID, container.Image, awc.input, aslw.StreamFlowAssistantMsg)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to get assistant provider", err)
+	}
+
+	msgChainID, err := assistantProvider.PrepareAgentChain(ctx)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to prepare assistant chain", err)
+	}
+
+	functionsBlob, err := json.Marshal(awc.functions)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to marshal functions", err)
+	}
+
+	logger = logger.WithField("msg_chain_id", msgChainID)
+	logger.Info("assistant provider prepared")
+
+	assistant, err = awc.db.UpdateAssistant(ctx, database.UpdateAssistantParams{
+		Title:              assistantProvider.Title(),
+		Model:              assistantProvider.Model(pconfig.OptionsTypePrimaryAgent),
+		Language:           assistantProvider.Language(),
+		ToolCallIDTemplate: assistantProvider.ToolCallIDTemplate(),
+		Functions:          functionsBlob,
+		TraceID:            database.StringToNullString(observation.TraceID()),
+		MsgchainID:         database.Int64ToNullInt64(&msgChainID),
+		ID:                 assistant.ID,
+	})
+	if err != nil {
+		logger.WithError(err).Error("failed to create assistant in DB")
+		return nil, fmt.Errorf("failed to create assistant in DB: %w", err)
+	}
+
+	workers, err := getFlowProviderWorkers(ctx, awc.flowID, &awc.flowProviderControllers)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to get flow provider workers", err)
+	}
+
+	assistantProvider.SetAgentLogProvider(workers.alw)
+	assistantProvider.SetMsgLogProvider(aslw)
+	assistantProvider.SetFlowWorker(awc.fw)
+
+	executor.SetImage(container.Image)
+	executor.SetEmbedder(assistantProvider.Embedder())
+	executor.SetScreenshotProvider(workers.sw)
+	executor.SetAgentLogProvider(workers.alw)
+	executor.SetMsgLogProvider(aslw)
+	executor.SetSearchLogProvider(workers.slw)
+	executor.SetTermLogProvider(workers.tlw)
+	executor.SetVectorStoreLogProvider(workers.vslw)
+	executor.SetToolCallLogProvider(workers.tclw)
+	executor.SetKnowledgeProvider(pub)
+	executor.SetGraphitiClient(awc.provs.GraphitiClient())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	ctx, _ = obs.Observer.NewObservation(ctx, langfuse.WithObservationTraceID(observation.TraceID()))
+	aw := &assistantWorker{
+		id:      assistant.ID,
+		flowID:  awc.flowID,
+		userID:  awc.userID,
+		chainID: msgChainID,
+		aslw:    aslw,
+		ap:      assistantProvider,
+		db:      awc.db,
+		wg:      &sync.WaitGroup{},
+		pub:     pub,
+		fw:      awc.fw,
+		ctx:     ctx,
+		cancel:  cancel,
+		runMX:   &sync.Mutex{},
+		runST:   func() {},
+		runWG:   &sync.WaitGroup{},
+		input:   make(chan assistantInput),
+		logger: logrus.WithFields(logrus.Fields{
+			"msg_chain_id": msgChainID,
+			"assistant_id": assistant.ID,
+			"flow_id":      awc.flowID,
+			"user_id":      awc.userID,
+			"trace_id":     observation.TraceID(),
+			"component":    "assistant",
+		}),
+	}
+
+	pub.AssistantCreated(ctx, assistant)
+
+	aw.wg.Add(1)
+	go aw.worker()
+
+	if err := aw.PutInput(ctx, awc.input, awc.useAgents, awc.resources); err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to run assistant worker", err)
+	}
+
+	assistantSpan.End(langfuse.WithSpanStatus("assistant worker started"))
+
+	return aw, nil
+}
+
+func LoadAssistantWorker(
+	ctx context.Context, assistant database.Assistant, awc assistantWorkerCtx,
+) (AssistantWorker, error) {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.LoadAssistantWorker")
+	defer span.End()
+
+	switch assistant.Status {
+	case database.AssistantStatusRunning, database.AssistantStatusWaiting:
+	default:
+		return nil, fmt.Errorf("assistant %d has status %s: loading aborted: %w", assistant.ID, assistant.Status, ErrNothingToLoad)
+	}
+
+	logger := logrus.WithContext(ctx).WithFields(logrus.Fields{
+		"assistant_id":  assistant.ID,
+		"flow_id":       awc.flowID,
+		"user_id":       awc.userID,
+		"msg_chain_id":  assistant.MsgchainID,
+		"provider_name": assistant.ModelProviderName,
+		"provider_type": assistant.ModelProviderType,
+	})
+
+	user, err := awc.db.GetUser(ctx, awc.userID)
+	if err != nil {
+		logger.WithError(err).Error("failed to get user")
+		return nil, fmt.Errorf("failed to get user %d: %w", awc.userID, err)
+	}
+
+	container, err := awc.db.GetFlowPrimaryContainer(ctx, awc.flowID)
+	if err != nil {
+		logger.WithError(err).Error("failed to get flow primary container")
+		return nil, fmt.Errorf("failed to get flow primary container: %w", err)
+	}
+
+	ctx, observation := obs.Observer.NewObservation(ctx,
+		langfuse.WithObservationTraceContext(
+			langfuse.WithTraceName(fmt.Sprintf("%d flow %d assistant worker", awc.flowID, assistant.ID)),
+			langfuse.WithTraceUserID(user.Mail),
+			langfuse.WithTraceTags([]string{"controller", "assistant"}),
+			langfuse.WithTraceSessionID(fmt.Sprintf("assistant-%d-flow-%d", assistant.ID, awc.flowID)),
+			langfuse.WithTraceMetadata(langfuse.Metadata{
+				"assistant_id":  assistant.ID,
+				"flow_id":       awc.flowID,
+				"user_id":       awc.userID,
+				"user_email":    user.Mail,
+				"user_name":     user.Name,
+				"user_hash":     user.Hash,
+				"user_role":     user.RoleName,
+				"provider_name": assistant.ModelProviderName,
+				"provider_type": assistant.ModelProviderType,
+			}),
+		),
+	)
+	assistantSpan := observation.Span(langfuse.WithSpanName("prepare assistant worker"))
+	ctx, _ = assistantSpan.Observation(ctx)
+
+	functions := &tools.Functions{}
+	if err := json.Unmarshal(assistant.Functions, functions); err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to unmarshal functions", err)
+	}
+
+	pub := awc.subs.NewFlowPublisher(awc.userID, awc.flowID)
+	aslw, err := awc.aslc.NewFlowAssistantLog(ctx, awc.flowID, assistant.ID, pub)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow assistant log worker", err)
+	}
+
+	prompter := awc.prompter
+	if prompter == nil {
+		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.userID, awc.flowID)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow tools executor", err)
+	}
+	assistantProvider, err := awc.provs.LoadAssistantProvider(ctx, provider.ProviderName(assistant.ModelProviderName),
+		prompter, executor, assistant.ID, awc.flowID, awc.userID, container.Image, assistant.Language, assistant.Title,
+		assistant.ToolCallIDTemplate, aslw.StreamFlowAssistantMsg)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to get assistant provider", err)
+	}
+
+	workers, err := getFlowProviderWorkers(ctx, awc.flowID, &awc.flowProviderControllers)
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to get flow provider workers", err)
+	}
+
+	assistantProvider.SetAgentLogProvider(workers.alw)
+	assistantProvider.SetMsgLogProvider(aslw)
+	assistantProvider.SetFlowWorker(awc.fw)
+
+	executor.SetImage(container.Image)
+	executor.SetEmbedder(assistantProvider.Embedder())
+	executor.SetScreenshotProvider(workers.sw)
+	executor.SetAgentLogProvider(workers.alw)
+	executor.SetMsgLogProvider(aslw)
+	executor.SetSearchLogProvider(workers.slw)
+	executor.SetTermLogProvider(workers.tlw)
+	executor.SetVectorStoreLogProvider(workers.vslw)
+	executor.SetToolCallLogProvider(workers.tclw)
+	executor.SetKnowledgeProvider(pub)
+
+	var msgChainID int64
+	pmsgChainID := database.NullInt64ToInt64(assistant.MsgchainID)
+	if pmsgChainID != nil {
+		msgChainID = *pmsgChainID
+		assistantProvider.SetMsgChainID(msgChainID)
+	} else {
+		return nil, fmt.Errorf("assistant %d has no msgchain id", assistant.ID)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	ctx, _ = obs.Observer.NewObservation(ctx, langfuse.WithObservationTraceID(observation.TraceID()))
+	aw := &assistantWorker{
+		id:      assistant.ID,
+		flowID:  awc.flowID,
+		userID:  awc.userID,
+		chainID: msgChainID,
+		aslw:    aslw,
+		ap:      assistantProvider,
+		db:      awc.db,
+		wg:      &sync.WaitGroup{},
+		pub:     pub,
+		fw:      awc.fw,
+		ctx:     ctx,
+		cancel:  cancel,
+		runMX:   &sync.Mutex{},
+		runST:   func() {},
+		runWG:   &sync.WaitGroup{},
+		input:   make(chan assistantInput),
+		logger: logrus.WithFields(logrus.Fields{
+			"msg_chain_id": msgChainID,
+			"assistant_id": assistant.ID,
+			"flow_id":      awc.flowID,
+			"user_id":      awc.userID,
+			"trace_id":     observation.TraceID(),
+			"component":    "assistant",
+		}),
+	}
+
+	assistant, err = awc.db.UpdateAssistantStatus(ctx, database.UpdateAssistantStatusParams{
+		Status: database.AssistantStatusWaiting,
+		ID:     assistant.ID,
+	})
+	if err != nil {
+		return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to update assistant status", err)
+	}
+
+	pub.AssistantUpdated(ctx, assistant)
+
+	aw.wg.Add(1)
+	go aw.worker()
+
+	assistantSpan.End(langfuse.WithSpanStatus("assistant worker started"))
+
+	return aw, nil
+}
+
+func (aw *assistantWorker) worker() {
+	defer aw.wg.Done()
+
+	perform := func(ctx context.Context, input string, useAgents bool) error {
+		aw.runWG.Add(1)
+		defer aw.runWG.Done()
+
+		_, err := aw.db.UpdateAssistantUseAgents(ctx, database.UpdateAssistantUseAgentsParams{
+			UseAgents: useAgents,
+			ID:        aw.id,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to update assistant use agents: %w", err)
+		}
+
+		if err := aw.SetStatus(ctx, database.AssistantStatusRunning); err != nil {
+			aw.logger.WithError(err).Error("failed to set assistant status to waiting")
+		}
+
+		defer func() {
+			if err := aw.SetStatus(ctx, database.AssistantStatusWaiting); err != nil {
+				aw.logger.WithError(err).Error("failed to set assistant status to waiting")
+			}
+		}()
+
+		_, err = aw.aslw.PutFlowAssistantMsg(ctx, database.MsglogTypeInput, "", input)
+		if err != nil {
+			return fmt.Errorf("failed to put input to flow assistant log: %w", err)
+		}
+
+		aw.runMX.Lock()
+		ctx, aw.runST = context.WithCancel(aw.ctx)
+		aw.runMX.Unlock()
+
+		if err := aw.ap.PutInputToAgentChain(ctx, input); err != nil {
+			return fmt.Errorf("failed to put input to agent chain: %w", err)
+		}
+
+		if err := aw.ap.PerformAgentChain(ctx); err != nil {
+			if errors.Is(err, context.Canceled) {
+				ctx = context.Background()
+			}
+			errChainConsistency := aw.ap.EnsureChainConsistency(ctx)
+			if errChainConsistency != nil {
+				err = errors.Join(err, errChainConsistency)
+			}
+			return fmt.Errorf("failed to perform agent chain: %w", err)
+		}
+
+		return nil
+	}
+
+	for {
+		select {
+		case <-aw.ctx.Done():
+			return
+		case ain := <-aw.input:
+			err := perform(aw.ctx, ain.input, ain.useAgents)
+			if err != nil {
+				aw.logger.WithError(err).Error("failed to perform assistant chain")
+			}
+			ain.done <- err
+		}
+	}
+}
+
+func (aw *assistantWorker) GetAssistantID() int64 {
+	return aw.id
+}
+
+func (aw *assistantWorker) GetUserID() int64 {
+	return aw.userID
+}
+
+func (aw *assistantWorker) GetFlowID() int64 {
+	return aw.flowID
+}
+
+func (aw *assistantWorker) GetTitle() string {
+	return aw.ap.Title()
+}
+
+func (aw *assistantWorker) GetStatus(ctx context.Context) (database.AssistantStatus, error) {
+	assistant, err := aw.db.GetAssistant(ctx, aw.id)
+	if err != nil {
+		return database.AssistantStatusFailed, err
+	}
+
+	return assistant.Status, nil
+}
+
+func (aw *assistantWorker) SetStatus(ctx context.Context, status database.AssistantStatus) error {
+	assistant, err := aw.db.UpdateAssistantStatus(ctx, database.UpdateAssistantStatusParams{
+		Status: status,
+		ID:     aw.id,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to update assistant %d flow %d status: %w", aw.id, aw.flowID, err)
+	}
+
+	aw.pub.AssistantUpdated(ctx, assistant)
+
+	return nil
+}
+
+func (aw *assistantWorker) PutInput(ctx context.Context, input string, useAgents bool, resources []database.UserResource) error {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.assistantWorker.PutInput")
+	defer span.End()
+
+	if aw.fw != nil {
+		if err := aw.fw.PutResources(ctx, resources); err != nil {
+			aw.logger.WithError(err).Warn("failed to copy resources before assistant input")
+		}
+	}
+
+	ain := assistantInput{input: input, useAgents: useAgents, done: make(chan error, 1)}
+	select {
+	case <-aw.ctx.Done():
+		close(ain.done)
+		return fmt.Errorf("assistant %d flow %d stopped: %w", aw.id, aw.flowID, aw.ctx.Err())
+	case <-ctx.Done():
+		close(ain.done)
+		return fmt.Errorf("assistant %d flow %d input processing timeout: %w", aw.id, aw.flowID, ctx.Err())
+	case aw.input <- ain:
+		timer := time.NewTimer(assistantInputTimeout)
+		defer timer.Stop()
+
+		select {
+		case err := <-ain.done:
+			return err // nil or error
+		case <-timer.C:
+			return nil // no early error
+		case <-aw.ctx.Done():
+			return fmt.Errorf("assistant %d flow %d stopped: %w", aw.id, aw.flowID, aw.ctx.Err())
+		case <-ctx.Done():
+			return fmt.Errorf("assistant %d flow %d input processing timeout: %w", aw.id, aw.flowID, ctx.Err())
+		}
+	}
+}
+
+func (aw *assistantWorker) Finish(ctx context.Context) error {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.assistantWorker.Finish")
+	defer span.End()
+
+	if err := aw.ctx.Err(); err != nil {
+		if errors.Is(err, context.Canceled) {
+			return nil
+		}
+		return fmt.Errorf("assistant %d flow %d stop failed: %w", aw.id, aw.flowID, err)
+	}
+
+	aw.cancel()
+	aw.wg.Wait()
+
+	if err := aw.SetStatus(ctx, database.AssistantStatusFinished); err != nil {
+		aw.logger.WithError(err).Error("failed to set assistant status to finished")
+	}
+
+	return nil
+}
+
+func (aw *assistantWorker) Stop(ctx context.Context) error {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.assistantWorker.Stop")
+	defer span.End()
+
+	aw.runST()
+	done := make(chan struct{})
+	timer := time.NewTimer(stopAssistantTimeout)
+	defer timer.Stop()
+
+	go func() {
+		aw.runWG.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-timer.C:
+		return fmt.Errorf("assistant stop timeout")
+	case <-done:
+		return nil
+	}
+}

backend/pkg/controller/context.go

@@ -0,0 +1,60 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+	"pentagi/pkg/observability/langfuse"
+	"pentagi/pkg/providers"
+	"pentagi/pkg/tools"
+
+	"github.com/sirupsen/logrus"
+)
+
+var ErrNothingToLoad = errors.New("nothing to load")
+
+type FlowContext struct {
+	DB database.Querier
+
+	UserID  int64
+	FlowID  int64
+	TraceID string
+
+	Executor  tools.FlowToolsExecutor
+	Provider  providers.FlowProvider
+	Publisher subscriptions.FlowPublisher
+
+	TermLog    FlowTermLogWorker
+	MsgLog     FlowMsgLogWorker
+	Screenshot FlowScreenshotWorker
+}
+
+type TaskContext struct {
+	TaskID    int64
+	TaskTitle string
+	TaskInput string
+
+	FlowContext
+}
+
+type SubtaskContext struct {
+	MsgChainID         int64
+	SubtaskID          int64
+	SubtaskTitle       string
+	SubtaskDescription string
+
+	TaskContext
+}
+
+func wrapErrorEndSpan(ctx context.Context, span langfuse.Span, msg string, err error) error {
+	logrus.WithContext(ctx).WithError(err).Error(msg)
+	err = fmt.Errorf("%s: %w", msg, err)
+	span.End(
+		langfuse.WithSpanStatus(err.Error()),
+		langfuse.WithSpanLevel(langfuse.ObservationLevelError),
+	)
+	return err
+}

backend/pkg/controller/flows.go

@@ -0,0 +1,391 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"sync"
+
+	"pentagi/pkg/config"
+	"pentagi/pkg/database"
+	"pentagi/pkg/docker"
+	"pentagi/pkg/graph/subscriptions"
+	"pentagi/pkg/providers"
+	"pentagi/pkg/providers/provider"
+	"pentagi/pkg/tools"
+
+	"github.com/sirupsen/logrus"
+)
+
+var (
+	ErrFlowNotFound       = fmt.Errorf("flow not found")
+	ErrFlowAlreadyStopped = fmt.Errorf("flow already stopped")
+)
+
+type FlowController interface {
+	CreateFlow(
+		ctx context.Context,
+		userID int64,
+		input string,
+		prvname provider.ProviderName,
+		prvtype provider.ProviderType,
+		functions *tools.Functions,
+		resources []database.UserResource,
+	) (FlowWorker, error)
+	CreateAssistant(
+		ctx context.Context,
+		userID int64,
+		flowID int64,
+		input string,
+		useAgents bool,
+		prvname provider.ProviderName,
+		prvtype provider.ProviderType,
+		functions *tools.Functions,
+		resources []database.UserResource,
+	) (AssistantWorker, error)
+	LoadFlows(ctx context.Context) error
+	ListFlows(ctx context.Context) []FlowWorker
+	GetFlow(ctx context.Context, flowID int64) (FlowWorker, error)
+	StopFlow(ctx context.Context, flowID int64) error
+	FinishFlow(ctx context.Context, flowID int64) error
+	RenameFlow(ctx context.Context, flowID int64, title string) error
+}
+
+type flowController struct {
+	db     database.Querier
+	mx     *sync.Mutex
+	cfg    *config.Config
+	flows  map[int64]FlowWorker
+	docker docker.DockerClient
+	provs  providers.ProviderController
+	subs   subscriptions.SubscriptionsController
+	alc    AgentLogController
+	mlc    MsgLogController
+	aslc   AssistantLogController
+	slc    SearchLogController
+	tlc    TermLogController
+	vslc   VectorStoreLogController
+	tclc   ToolCallLogController
+	sc     ScreenshotController
+}
+
+func NewFlowController(
+	db database.Querier,
+	cfg *config.Config,
+	docker docker.DockerClient,
+	provs providers.ProviderController,
+	subs subscriptions.SubscriptionsController,
+) FlowController {
+	return &flowController{
+		db:     db,
+		mx:     &sync.Mutex{},
+		cfg:    cfg,
+		flows:  make(map[int64]FlowWorker),
+		docker: docker,
+		provs:  provs,
+		subs:   subs,
+		alc:    NewAgentLogController(db),
+		mlc:    NewMsgLogController(db),
+		aslc:   NewAssistantLogController(db),
+		slc:    NewSearchLogController(db),
+		tlc:    NewTermLogController(db),
+		vslc:   NewVectorStoreLogController(db),
+		tclc:   NewToolCallLogController(db),
+		sc:     NewScreenshotController(db),
+	}
+}
+
+func (fc *flowController) LoadFlows(ctx context.Context) error {
+	flows, err := fc.db.GetFlows(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to load flows: %w", err)
+	}
+
+	for _, flow := range flows {
+		fw, err := LoadFlowWorker(ctx, flow, flowWorkerCtx{
+			db:     fc.db,
+			cfg:    fc.cfg,
+			docker: fc.docker,
+			provs:  fc.provs,
+			subs:   fc.subs,
+			flowProviderControllers: flowProviderControllers{
+				mlc:  fc.mlc,
+				aslc: fc.aslc,
+				alc:  fc.alc,
+				slc:  fc.slc,
+				tlc:  fc.tlc,
+				vslc: fc.vslc,
+				tclc: fc.tclc,
+				sc:   fc.sc,
+			},
+		})
+		if err != nil {
+			if errors.Is(err, ErrNothingToLoad) {
+				continue
+			}
+
+			logrus.WithContext(ctx).WithError(err).Errorf("failed to load flow %d", flow.ID)
+			continue
+		}
+
+		fc.flows[flow.ID] = fw
+	}
+
+	return nil
+}
+
+func (fc *flowController) CreateFlow(
+	ctx context.Context,
+	userID int64,
+	input string,
+	prvname provider.ProviderName,
+	prvtype provider.ProviderType,
+	functions *tools.Functions,
+	resources []database.UserResource,
+) (FlowWorker, error) {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	fw, err := NewFlowWorker(ctx, newFlowWorkerCtx{
+		userID:    userID,
+		input:     input,
+		prvname:   prvname,
+		prvtype:   prvtype,
+		functions: functions,
+		resources: resources,
+		flowWorkerCtx: flowWorkerCtx{
+			db:     fc.db,
+			cfg:    fc.cfg,
+			docker: fc.docker,
+			provs:  fc.provs,
+			subs:   fc.subs,
+			flowProviderControllers: flowProviderControllers{
+				mlc:  fc.mlc,
+				aslc: fc.aslc,
+				alc:  fc.alc,
+				slc:  fc.slc,
+				tlc:  fc.tlc,
+				vslc: fc.vslc,
+				tclc: fc.tclc,
+				sc:   fc.sc,
+			},
+		},
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to create flow worker: %w", err)
+	}
+
+	fc.flows[fw.GetFlowID()] = fw
+
+	return fw, nil
+}
+
+func (fc *flowController) CreateAssistant(
+	ctx context.Context,
+	userID int64,
+	flowID int64,
+	input string,
+	useAgents bool,
+	prvname provider.ProviderName,
+	prvtype provider.ProviderType,
+	functions *tools.Functions,
+	resources []database.UserResource,
+) (AssistantWorker, error) {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	var (
+		fw  FlowWorker
+		ok  bool
+		err error
+	)
+
+	flowWorkerCtx := flowWorkerCtx{
+		db:     fc.db,
+		cfg:    fc.cfg,
+		docker: fc.docker,
+		provs:  fc.provs,
+		subs:   fc.subs,
+		flowProviderControllers: flowProviderControllers{
+			mlc:  fc.mlc,
+			aslc: fc.aslc,
+			alc:  fc.alc,
+			slc:  fc.slc,
+			tlc:  fc.tlc,
+			vslc: fc.vslc,
+			tclc: fc.tclc,
+			sc:   fc.sc,
+		},
+	}
+
+	newFlow := func() error {
+		fw, err = NewFlowWorker(ctx, newFlowWorkerCtx{
+			userID:        userID,
+			input:         input,
+			dryRun:        true,
+			prvname:       prvname,
+			prvtype:       prvtype,
+			functions:     functions,
+			flowWorkerCtx: flowWorkerCtx,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to create flow worker: %w", err)
+		}
+
+		fc.flows[fw.GetFlowID()] = fw
+		flowID = fw.GetFlowID()
+		fw.SetStatus(ctx, database.FlowStatusWaiting)
+
+		return nil
+	}
+
+	loadFlow := func() error {
+		flow, err := fc.db.UpdateFlowStatus(ctx, database.UpdateFlowStatusParams{
+			ID:     flowID,
+			Status: database.FlowStatusWaiting,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to renew flow %d status: %w", flowID, err)
+		}
+
+		fw, err = LoadFlowWorker(ctx, flow, flowWorkerCtx)
+		if err != nil {
+			return fmt.Errorf("failed to load flow %d: %w", flowID, err)
+		}
+
+		fc.flows[flowID] = fw
+
+		return nil
+	}
+
+	if flowID == 0 {
+		if err := newFlow(); err != nil {
+			return nil, err
+		}
+	} else if fw, ok = fc.flows[flowID]; ok {
+		status, err := fw.GetStatus(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get flow %d status: %w", flowID, err)
+		}
+
+		switch status {
+		case database.FlowStatusCreated:
+			return nil, fmt.Errorf("flow %d is not completed", flowID)
+		case database.FlowStatusFinished, database.FlowStatusFailed:
+			if err := loadFlow(); err != nil {
+				return nil, err
+			}
+		case database.FlowStatusRunning, database.FlowStatusWaiting:
+			break
+		default:
+			return nil, fmt.Errorf("flow %d is in unknown status: %s", flowID, status)
+		}
+	} else {
+		if err := loadFlow(); err != nil {
+			return nil, err
+		}
+	}
+
+	if fw == nil { // just double check, this should never happen
+		return nil, fmt.Errorf("unexpected error: flow %d not found", flowID)
+	}
+
+	aw, err := NewAssistantWorker(ctx, newAssistantWorkerCtx{
+		userID:        userID,
+		flowID:        flowID,
+		input:         input,
+		prvname:       prvname,
+		prvtype:       prvtype,
+		useAgents:     useAgents,
+		functions:     functions,
+		resources:     resources,
+		fw:            fw,
+		flowWorkerCtx: flowWorkerCtx,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to create assistant: %w", err)
+	}
+
+	if err = fw.AddAssistant(ctx, aw); err != nil {
+		return nil, fmt.Errorf("failed to add assistant to flow: %w", err)
+	}
+
+	return aw, nil
+}
+
+func (fc *flowController) ListFlows(ctx context.Context) []FlowWorker {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	flows := make([]FlowWorker, 0)
+	for _, flow := range fc.flows {
+		flows = append(flows, flow)
+	}
+
+	sort.Slice(flows, func(i, j int) bool {
+		return flows[i].GetFlowID() < flows[j].GetFlowID()
+	})
+
+	return flows
+}
+
+func (fc *flowController) GetFlow(ctx context.Context, flowID int64) (FlowWorker, error) {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	flow, ok := fc.flows[flowID]
+	if !ok {
+		return nil, ErrFlowNotFound
+	}
+
+	return flow, nil
+}
+
+func (fc *flowController) StopFlow(ctx context.Context, flowID int64) error {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	flow, ok := fc.flows[flowID]
+	if !ok {
+		return ErrFlowNotFound
+	}
+
+	err := flow.Stop(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to stop flow %d: %w", flowID, err)
+	}
+
+	return nil
+}
+
+func (fc *flowController) FinishFlow(ctx context.Context, flowID int64) error {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	flow, ok := fc.flows[flowID]
+	if !ok {
+		return ErrFlowNotFound
+	}
+
+	err := flow.Finish(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to finish flow %d: %w", flowID, err)
+	}
+
+	delete(fc.flows, flowID)
+
+	return nil
+}
+
+func (fc *flowController) RenameFlow(ctx context.Context, flowID int64, title string) error {
+	fc.mx.Lock()
+	defer fc.mx.Unlock()
+
+	flow, ok := fc.flows[flowID]
+	if !ok {
+		return ErrFlowNotFound
+	}
+
+	return flow.Rename(ctx, title)
+}

backend/pkg/controller/msglog.go

@@ -0,0 +1,243 @@
+package controller
+
+import (
+	"context"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+const defaultMaxMessageLength = 2048
+
+type FlowMsgLogWorker interface {
+	PutMsg(
+		ctx context.Context,
+		msgType database.MsglogType,
+		taskID, subtaskID *int64,
+		streamID int64,
+		thinking, msg string,
+	) (int64, error)
+	PutFlowMsg(
+		ctx context.Context,
+		msgType database.MsglogType,
+		thinking, msg string,
+	) (int64, error)
+	PutFlowMsgResult(
+		ctx context.Context,
+		msgType database.MsglogType,
+		thinking, msg, result string,
+		resultFormat database.MsglogResultFormat,
+	) (int64, error)
+	PutTaskMsg(
+		ctx context.Context,
+		msgType database.MsglogType,
+		taskID int64,
+		thinking, msg string,
+	) (int64, error)
+	PutTaskMsgResult(
+		ctx context.Context,
+		msgType database.MsglogType,
+		taskID int64,
+		thinking, msg, result string,
+		resultFormat database.MsglogResultFormat,
+	) (int64, error)
+	PutSubtaskMsg(
+		ctx context.Context,
+		msgType database.MsglogType,
+		taskID, subtaskID int64,
+		thinking, msg string,
+	) (int64, error)
+	PutSubtaskMsgResult(
+		ctx context.Context,
+		msgType database.MsglogType,
+		taskID, subtaskID int64,
+		thinking, msg, result string,
+		resultFormat database.MsglogResultFormat,
+	) (int64, error)
+	UpdateMsgResult(
+		ctx context.Context,
+		msgID, streamID int64,
+		result string,
+		resultFormat database.MsglogResultFormat,
+	) error
+}
+
+type flowMsgLogWorker struct {
+	db     database.Querier
+	mx     *sync.Mutex
+	flowID int64
+	pub    subscriptions.FlowPublisher
+}
+
+func NewFlowMsgLogWorker(db database.Querier, flowID int64, pub subscriptions.FlowPublisher) FlowMsgLogWorker {
+	return &flowMsgLogWorker{
+		db:     db,
+		mx:     &sync.Mutex{},
+		flowID: flowID,
+		pub:    pub,
+	}
+}
+
+func (mlw *flowMsgLogWorker) PutMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID *int64,
+	streamID int64, // unsupported for now
+	thinking, msg string,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsg(ctx, msgType, taskID, subtaskID, thinking, msg)
+}
+
+func (mlw *flowMsgLogWorker) PutFlowMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	thinking, msg string,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsg(ctx, msgType, nil, nil, thinking, msg)
+}
+
+func (mlw *flowMsgLogWorker) PutFlowMsgResult(
+	ctx context.Context,
+	msgType database.MsglogType,
+	thinking, msg, result string,
+	resultFormat database.MsglogResultFormat,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsgResult(ctx, msgType, nil, nil, thinking, msg, result, resultFormat)
+}
+
+func (mlw *flowMsgLogWorker) PutTaskMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID int64,
+	thinking, msg string,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsg(ctx, msgType, &taskID, nil, thinking, msg)
+}
+
+func (mlw *flowMsgLogWorker) PutTaskMsgResult(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID int64,
+	thinking, msg, result string,
+	resultFormat database.MsglogResultFormat,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsgResult(ctx, msgType, &taskID, nil, thinking, msg, result, resultFormat)
+}
+
+func (mlw *flowMsgLogWorker) PutSubtaskMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID int64,
+	thinking, msg string,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsg(ctx, msgType, &taskID, &subtaskID, thinking, msg)
+}
+
+func (mlw *flowMsgLogWorker) PutSubtaskMsgResult(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID int64,
+	thinking, msg, result string,
+	resultFormat database.MsglogResultFormat,
+) (int64, error) {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	return mlw.putMsgResult(ctx, msgType, &taskID, &subtaskID, thinking, msg, result, resultFormat)
+}
+
+func (mlw *flowMsgLogWorker) UpdateMsgResult(
+	ctx context.Context,
+	msgID int64,
+	streamID int64, // unsupported for now
+	result string,
+	resultFormat database.MsglogResultFormat,
+) error {
+	mlw.mx.Lock()
+	defer mlw.mx.Unlock()
+
+	msgLog, err := mlw.db.UpdateMsgLogResult(ctx, database.UpdateMsgLogResultParams{
+		Result:       database.SanitizeUTF8(result),
+		ResultFormat: resultFormat,
+		ID:           msgID,
+	})
+	if err != nil {
+		return err
+	}
+
+	mlw.pub.MessageLogUpdated(ctx, msgLog)
+
+	return nil
+}
+
+func (mlw *flowMsgLogWorker) putMsg(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID *int64,
+	thinking, msg string,
+) (int64, error) {
+	msgLog, err := mlw.db.CreateMsgLog(ctx, database.CreateMsgLogParams{
+		Type:      msgType,
+		Message:   database.SanitizeUTF8(msg),
+		Thinking:  database.StringToNullString(database.SanitizeUTF8(thinking)),
+		FlowID:    mlw.flowID,
+		TaskID:    database.Int64ToNullInt64(taskID),
+		SubtaskID: database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	mlw.pub.MessageLogAdded(ctx, msgLog)
+
+	return msgLog.ID, nil
+}
+
+func (mlw *flowMsgLogWorker) putMsgResult(
+	ctx context.Context,
+	msgType database.MsglogType,
+	taskID, subtaskID *int64,
+	thinking, msg, result string,
+	resultFormat database.MsglogResultFormat,
+) (int64, error) {
+	if len(msg) > defaultMaxMessageLength {
+		msg = msg[:defaultMaxMessageLength] + "..."
+	}
+
+	msgLog, err := mlw.db.CreateResultMsgLog(ctx, database.CreateResultMsgLogParams{
+		Type:         msgType,
+		Message:      database.SanitizeUTF8(msg),
+		Thinking:     database.StringToNullString(database.SanitizeUTF8(thinking)),
+		Result:       database.SanitizeUTF8(result),
+		ResultFormat: resultFormat,
+		FlowID:       mlw.flowID,
+		TaskID:       database.Int64ToNullInt64(taskID),
+		SubtaskID:    database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, err
+	}
+
+	mlw.pub.MessageLogAdded(ctx, msgLog)
+
+	return msgLog.ID, nil
+}

backend/pkg/controller/msglogs.go

@@ -0,0 +1,68 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type MsgLogController interface {
+	NewFlowMsgLog(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowMsgLogWorker, error)
+	ListFlowsMsgLog(ctx context.Context) ([]FlowMsgLogWorker, error)
+	GetFlowMsgLog(ctx context.Context, flowID int64) (FlowMsgLogWorker, error)
+}
+
+type msgLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowMsgLogWorker
+}
+
+func NewMsgLogController(db database.Querier) MsgLogController {
+	return &msgLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowMsgLogWorker),
+	}
+}
+
+func (mlc *msgLogController) NewFlowMsgLog(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowMsgLogWorker, error) {
+	mlc.mx.Lock()
+	defer mlc.mx.Unlock()
+
+	flw := NewFlowMsgLogWorker(mlc.db, flowID, pub)
+	mlc.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (mlc *msgLogController) ListFlowsMsgLog(ctx context.Context) ([]FlowMsgLogWorker, error) {
+	mlc.mx.Lock()
+	defer mlc.mx.Unlock()
+
+	flows := make([]FlowMsgLogWorker, 0, len(mlc.flows))
+	for _, flw := range mlc.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (mlc *msgLogController) GetFlowMsgLog(ctx context.Context, flowID int64) (FlowMsgLogWorker, error) {
+	mlc.mx.Lock()
+	defer mlc.mx.Unlock()
+
+	flw, ok := mlc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}

backend/pkg/controller/prompter.go

@@ -0,0 +1,50 @@
+package controller
+
+import (
+	"context"
+	"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)
+	}
+
+	defaults, err := templates.LoadDefaultPromptsMap()
+	if err != nil {
+		return nil, fmt.Errorf("failed to load default templates: %w", err)
+	}
+
+	return buildUserPrompter(defaults, userPrompts), nil
+}
+
+// buildUserPrompter is the pure merge step extracted from newUserPrompter so
+// it can be unit-tested without a database fake or filesystem access. It
+// mutates the supplied defaults map by overlaying each non-empty user
+// override on top, then returns a Prompter backed by that map. Callers must
+// pass a fresh map (e.g., from templates.LoadDefaultPromptsMap) so the
+// embedded defaults are not modified.
+func buildUserPrompter(defaults templates.PromptsMap, userPrompts []database.Prompt) templates.Prompter {
+	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
+		}
+		defaults[templates.PromptType(p.Type)] = p.Prompt
+	}
+
+	return templates.NewFlowPrompter(defaults)
+}

backend/pkg/controller/prompter_test.go

@@ -0,0 +1,209 @@
+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
+}
+
+// loadDefaults returns a fresh PromptsMap of the embedded default templates so
+// each test mutates its own map without leaking state between cases.
+func loadDefaults(t *testing.T) templates.PromptsMap {
+	t.Helper()
+	defaults, err := templates.LoadDefaultPromptsMap()
+	if err != nil {
+		t.Fatalf("LoadDefaultPromptsMap error: %v", err)
+	}
+	return defaults
+}
+
+func TestBuildUserPrompter_NoUserPrompts(t *testing.T) {
+	prompter := buildUserPrompter(loadDefaults(t), nil)
+
+	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 := buildUserPrompter(loadDefaults(t), userPrompts)
+
+	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 := buildUserPrompter(loadDefaults(t), userPrompts)
+
+	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 := buildUserPrompter(loadDefaults(t), userPrompts)
+
+	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)
+	}
+}

backend/pkg/controller/screenshot.go

@@ -0,0 +1,62 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type FlowScreenshotWorker interface {
+	PutScreenshot(ctx context.Context, name, url string, taskID, subtaskID *int64) (int64, error)
+	GetScreenshot(ctx context.Context, screenshotID int64) (database.Screenshot, error)
+}
+
+type flowScreenshotWorker struct {
+	db         database.Querier
+	mx         *sync.Mutex
+	flowID     int64
+	containers map[int64]struct{}
+	pub        subscriptions.FlowPublisher
+}
+
+func NewFlowScreenshotWorker(db database.Querier, flowID int64, pub subscriptions.FlowPublisher) FlowScreenshotWorker {
+	return &flowScreenshotWorker{
+		db:         db,
+		mx:         &sync.Mutex{},
+		flowID:     flowID,
+		containers: make(map[int64]struct{}),
+		pub:        pub,
+	}
+}
+
+func (sw *flowScreenshotWorker) PutScreenshot(ctx context.Context, name, url string, taskID, subtaskID *int64) (int64, error) {
+	sw.mx.Lock()
+	defer sw.mx.Unlock()
+
+	screenshot, err := sw.db.CreateScreenshot(ctx, database.CreateScreenshotParams{
+		Name:      database.SanitizeUTF8(name),
+		Url:       database.SanitizeUTF8(url),
+		FlowID:    sw.flowID,
+		TaskID:    database.Int64ToNullInt64(taskID),
+		SubtaskID: database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to create screenshot: %w", err)
+	}
+
+	sw.pub.ScreenshotAdded(ctx, screenshot)
+
+	return screenshot.ID, nil
+}
+
+func (sw *flowScreenshotWorker) GetScreenshot(ctx context.Context, screenshotID int64) (database.Screenshot, error) {
+	screenshot, err := sw.db.GetScreenshot(ctx, screenshotID)
+	if err != nil {
+		return database.Screenshot{}, fmt.Errorf("failed to get screenshot: %w", err)
+	}
+
+	return screenshot, nil
+}

backend/pkg/controller/screenshots.go

@@ -0,0 +1,68 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type ScreenshotController interface {
+	NewFlowScreenshot(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowScreenshotWorker, error)
+	ListFlowsScreenshot(ctx context.Context) ([]FlowScreenshotWorker, error)
+	GetFlowScreenshot(ctx context.Context, flowID int64) (FlowScreenshotWorker, error)
+}
+
+type screenshotController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowScreenshotWorker
+}
+
+func NewScreenshotController(db database.Querier) ScreenshotController {
+	return &screenshotController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowScreenshotWorker),
+	}
+}
+
+func (sc *screenshotController) NewFlowScreenshot(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowScreenshotWorker, error) {
+	sc.mx.Lock()
+	defer sc.mx.Unlock()
+
+	flw := NewFlowScreenshotWorker(sc.db, flowID, pub)
+	sc.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (sc *screenshotController) ListFlowsScreenshot(ctx context.Context) ([]FlowScreenshotWorker, error) {
+	sc.mx.Lock()
+	defer sc.mx.Unlock()
+
+	flows := make([]FlowScreenshotWorker, 0, len(sc.flows))
+	for _, flw := range sc.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (sc *screenshotController) GetFlowScreenshot(ctx context.Context, flowID int64) (FlowScreenshotWorker, error) {
+	sc.mx.Lock()
+	defer sc.mx.Unlock()
+
+	flw, ok := sc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}

backend/pkg/controller/slog.go

@@ -0,0 +1,85 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type FlowSearchLogWorker interface {
+	PutLog(
+		ctx context.Context,
+		initiator database.MsgchainType,
+		executor database.MsgchainType,
+		engine database.SearchengineType,
+		query string,
+		result string,
+		taskID *int64,
+		subtaskID *int64,
+	) (int64, error)
+	GetLog(ctx context.Context, msgID int64) (database.Searchlog, error)
+}
+
+type flowSearchLogWorker struct {
+	db         database.Querier
+	mx         *sync.Mutex
+	flowID     int64
+	containers map[int64]struct{}
+	pub        subscriptions.FlowPublisher
+}
+
+func NewFlowSearchLogWorker(db database.Querier, flowID int64, pub subscriptions.FlowPublisher) FlowSearchLogWorker {
+	return &flowSearchLogWorker{
+		db:     db,
+		mx:     &sync.Mutex{},
+		flowID: flowID,
+		pub:    pub,
+	}
+}
+
+func (slw *flowSearchLogWorker) PutLog(
+	ctx context.Context,
+	initiator database.MsgchainType,
+	executor database.MsgchainType,
+	engine database.SearchengineType,
+	query string,
+	result string,
+	taskID *int64,
+	subtaskID *int64,
+) (int64, error) {
+	slw.mx.Lock()
+	defer slw.mx.Unlock()
+
+	slLog, err := slw.db.CreateSearchLog(ctx, database.CreateSearchLogParams{
+		Initiator: initiator,
+		Executor:  executor,
+		Engine:    engine,
+		Query:     query,
+		Result:    result,
+		FlowID:    slw.flowID,
+		TaskID:    database.Int64ToNullInt64(taskID),
+		SubtaskID: database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to create search log: %w", err)
+	}
+
+	slw.pub.SearchLogAdded(ctx, slLog)
+
+	return slLog.ID, nil
+}
+
+func (slw *flowSearchLogWorker) GetLog(ctx context.Context, msgID int64) (database.Searchlog, error) {
+	msg, err := slw.db.GetFlowSearchLog(ctx, database.GetFlowSearchLogParams{
+		ID:     msgID,
+		FlowID: slw.flowID,
+	})
+	if err != nil {
+		return database.Searchlog{}, fmt.Errorf("failed to get search log: %w", err)
+	}
+
+	return msg, nil
+}

backend/pkg/controller/slogs.go

@@ -0,0 +1,71 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type SearchLogController interface {
+	NewFlowSearchLog(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowSearchLogWorker, error)
+	ListFlowsSearchLog(ctx context.Context) ([]FlowSearchLogWorker, error)
+	GetFlowSearchLog(ctx context.Context, flowID int64) (FlowSearchLogWorker, error)
+}
+
+type searchLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowSearchLogWorker
+}
+
+func NewSearchLogController(db database.Querier) SearchLogController {
+	return &searchLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowSearchLogWorker),
+	}
+}
+
+func (slc *searchLogController) NewFlowSearchLog(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowSearchLogWorker, error) {
+	slc.mx.Lock()
+	defer slc.mx.Unlock()
+
+	flw := NewFlowSearchLogWorker(slc.db, flowID, pub)
+	slc.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (slc *searchLogController) ListFlowsSearchLog(ctx context.Context) ([]FlowSearchLogWorker, error) {
+	slc.mx.Lock()
+	defer slc.mx.Unlock()
+
+	flows := make([]FlowSearchLogWorker, 0, len(slc.flows))
+	for _, flw := range slc.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (slc *searchLogController) GetFlowSearchLog(
+	ctx context.Context,
+	flowID int64,
+) (FlowSearchLogWorker, error) {
+	slc.mx.Lock()
+	defer slc.mx.Unlock()
+
+	flw, ok := slc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}

backend/pkg/controller/subtask.go

@@ -0,0 +1,371 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"pentagi/pkg/database"
+	obs "pentagi/pkg/observability"
+	"pentagi/pkg/providers"
+
+	"github.com/sirupsen/logrus"
+)
+
+type TaskUpdater interface {
+	SetStatus(ctx context.Context, status database.TaskStatus) error
+}
+
+type SubtaskWorker interface {
+	GetMsgChainID() int64
+	GetSubtaskID() int64
+	GetTaskID() int64
+	GetFlowID() int64
+	GetUserID() int64
+	GetTitle() string
+	GetDescription() string
+	IsCompleted() bool
+	IsWaiting() bool
+	GetStatus(ctx context.Context) (database.SubtaskStatus, error)
+	SetStatus(ctx context.Context, status database.SubtaskStatus) error
+	GetResult(ctx context.Context) (string, error)
+	SetResult(ctx context.Context, result string) error
+	PutInput(ctx context.Context, input string) error
+	Run(ctx context.Context) error
+	Finish(ctx context.Context) error
+}
+
+type subtaskWorker struct {
+	mx         *sync.RWMutex
+	subtaskCtx *SubtaskContext
+	updater    TaskUpdater
+	completed  bool
+	waiting    bool
+}
+
+func NewSubtaskWorker(
+	ctx context.Context,
+	taskCtx *TaskContext,
+	id int64,
+	title,
+	description string,
+	updater TaskUpdater,
+) (SubtaskWorker, error) {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.NewSubtaskWorker")
+	defer span.End()
+
+	msgChainID, err := taskCtx.Provider.PrepareAgentChain(ctx, taskCtx.TaskID, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to prepare primary agent chain for subtask %d: %w", id, err)
+	}
+
+	return &subtaskWorker{
+		mx: &sync.RWMutex{},
+		subtaskCtx: &SubtaskContext{
+			MsgChainID:         msgChainID,
+			SubtaskID:          id,
+			SubtaskTitle:       title,
+			SubtaskDescription: description,
+			TaskContext:        *taskCtx,
+		},
+		updater:   updater,
+		completed: false,
+		waiting:   false,
+	}, nil
+}
+
+func LoadSubtaskWorker(
+	ctx context.Context,
+	subtask database.Subtask,
+	taskCtx *TaskContext,
+	updater TaskUpdater,
+) (SubtaskWorker, error) {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.LoadSubtaskWorker")
+	defer span.End()
+
+	var completed, waiting bool
+	switch subtask.Status {
+	case database.SubtaskStatusFinished, database.SubtaskStatusFailed:
+		completed = true
+	case database.SubtaskStatusWaiting:
+		waiting = true
+	case database.SubtaskStatusRunning:
+		var err error
+		// if subtask is running, it means that it was not finished by previous run
+		// so we need to set it to created and continue from the beginning
+		subtask, err = taskCtx.DB.UpdateSubtaskStatus(ctx, database.UpdateSubtaskStatusParams{
+			Status: database.SubtaskStatusCreated,
+			ID:     subtask.ID,
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to update subtask %d status to created: %w", subtask.ID, err)
+		}
+	case database.SubtaskStatusCreated:
+		return nil, fmt.Errorf("subtask %d has created yet: %w", subtask.ID, ErrNothingToLoad)
+	default:
+		return nil, fmt.Errorf("unexpected subtask status: %s", subtask.Status)
+	}
+
+	msgChains, err := taskCtx.DB.GetSubtaskPrimaryMsgChains(ctx, database.Int64ToNullInt64(&subtask.ID))
+	if err != nil {
+		return nil, fmt.Errorf("failed to get subtask primary msg chains for subtask %d: %w", subtask.ID, err)
+	}
+
+	if len(msgChains) == 0 {
+		return nil, fmt.Errorf("subtask %d has no msg chains: %w", subtask.ID, ErrNothingToLoad)
+	}
+
+	return &subtaskWorker{
+		mx: &sync.RWMutex{},
+		subtaskCtx: &SubtaskContext{
+			MsgChainID:         msgChains[0].ID,
+			SubtaskID:          subtask.ID,
+			SubtaskTitle:       subtask.Title,
+			SubtaskDescription: subtask.Description,
+			TaskContext:        *taskCtx,
+		},
+		updater:   updater,
+		completed: completed,
+		waiting:   waiting,
+	}, nil
+}
+
+func (stw *subtaskWorker) GetMsgChainID() int64 {
+	return stw.subtaskCtx.MsgChainID
+}
+
+func (stw *subtaskWorker) GetSubtaskID() int64 {
+	return stw.subtaskCtx.SubtaskID
+}
+
+func (stw *subtaskWorker) GetTaskID() int64 {
+	return stw.subtaskCtx.TaskID
+}
+
+func (stw *subtaskWorker) GetFlowID() int64 {
+	return stw.subtaskCtx.FlowID
+}
+
+func (stw *subtaskWorker) GetUserID() int64 {
+	return stw.subtaskCtx.UserID
+}
+
+func (stw *subtaskWorker) GetTitle() string {
+	return stw.subtaskCtx.SubtaskTitle
+}
+
+func (stw *subtaskWorker) GetDescription() string {
+	return stw.subtaskCtx.SubtaskDescription
+}
+
+func (stw *subtaskWorker) IsCompleted() bool {
+	stw.mx.RLock()
+	defer stw.mx.RUnlock()
+
+	return stw.completed
+}
+
+func (stw *subtaskWorker) IsWaiting() bool {
+	stw.mx.RLock()
+	defer stw.mx.RUnlock()
+
+	return stw.waiting
+}
+
+func (stw *subtaskWorker) GetStatus(ctx context.Context) (database.SubtaskStatus, error) {
+	subtask, err := stw.subtaskCtx.DB.GetSubtask(ctx, stw.subtaskCtx.SubtaskID)
+	if err != nil {
+		return database.SubtaskStatusFailed, err
+	}
+
+	return subtask.Status, nil
+}
+
+func (stw *subtaskWorker) SetStatus(ctx context.Context, status database.SubtaskStatus) error {
+	_, err := stw.subtaskCtx.DB.UpdateSubtaskStatus(ctx, database.UpdateSubtaskStatusParams{
+		Status: status,
+		ID:     stw.subtaskCtx.SubtaskID,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to set subtask %d status: %w", stw.subtaskCtx.SubtaskID, err)
+	}
+
+	stw.mx.Lock()
+	defer stw.mx.Unlock()
+
+	switch status {
+	case database.SubtaskStatusRunning:
+		stw.completed = false
+		stw.waiting = false
+		err = stw.updater.SetStatus(ctx, database.TaskStatusRunning)
+	case database.SubtaskStatusWaiting:
+		stw.completed = false
+		stw.waiting = true
+		err = stw.updater.SetStatus(ctx, database.TaskStatusWaiting)
+	case database.SubtaskStatusFinished, database.SubtaskStatusFailed:
+		stw.completed = true
+		stw.waiting = false
+		// statuses Finished and Failed will be produced by stack from Run function call
+	default:
+		// status Created is not possible to set by this call
+		return fmt.Errorf("unsupported subtask status: %s", status)
+	}
+	if err != nil {
+		return fmt.Errorf("failed to set task status in back propagation: %w", err)
+	}
+
+	return nil
+}
+
+func (stw *subtaskWorker) GetResult(ctx context.Context) (string, error) {
+	subtask, err := stw.subtaskCtx.DB.GetSubtask(ctx, stw.subtaskCtx.SubtaskID)
+	if err != nil {
+		return "", err
+	}
+
+	return subtask.Result, nil
+}
+
+func (stw *subtaskWorker) SetResult(ctx context.Context, result string) error {
+	_, err := stw.subtaskCtx.DB.UpdateSubtaskResult(ctx, database.UpdateSubtaskResultParams{
+		Result: result,
+		ID:     stw.subtaskCtx.SubtaskID,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to set subtask %d result: %w", stw.subtaskCtx.SubtaskID, err)
+	}
+
+	return nil
+}
+
+func (stw *subtaskWorker) PutInput(ctx context.Context, input string) error {
+	if stw.IsCompleted() {
+		return fmt.Errorf("subtask has already completed")
+	}
+
+	if !stw.IsWaiting() {
+		return fmt.Errorf("subtask is not waiting, run first")
+	}
+
+	err := stw.subtaskCtx.Provider.PutInputToAgentChain(ctx, stw.subtaskCtx.MsgChainID, input)
+	if err != nil {
+		return fmt.Errorf("failed to put input for subtask %d: %w", stw.subtaskCtx.SubtaskID, err)
+	}
+
+	_, err = stw.subtaskCtx.MsgLog.PutSubtaskMsg(
+		ctx,
+		database.MsglogTypeInput,
+		stw.subtaskCtx.TaskID,
+		stw.subtaskCtx.SubtaskID,
+		"", // thinking is empty because this is input
+		input,
+	)
+	if err != nil {
+		return fmt.Errorf("failed to put input for subtask %d: %w", stw.subtaskCtx.SubtaskID, err)
+	}
+
+	stw.mx.Lock()
+	defer stw.mx.Unlock()
+
+	stw.waiting = false
+
+	return nil
+}
+
+func (stw *subtaskWorker) Run(ctx context.Context) error {
+	if stw.IsCompleted() {
+		return fmt.Errorf("subtask has already completed")
+	}
+
+	if stw.IsWaiting() {
+		return fmt.Errorf("subtask is waiting, put input first")
+	}
+
+	if err := stw.SetStatus(ctx, database.SubtaskStatusRunning); err != nil {
+		stw.handleInterrupting(err)
+		return err
+	}
+
+	var (
+		taskID     = stw.subtaskCtx.TaskID
+		subtaskID  = stw.subtaskCtx.SubtaskID
+		msgChainID = stw.subtaskCtx.MsgChainID
+	)
+
+	if err := stw.subtaskCtx.Provider.EnsureChainConsistency(ctx, msgChainID); err != nil {
+		stw.handleInterrupting(err)
+		return fmt.Errorf("failed to ensure chain consistency for subtask %d: %w", subtaskID, err)
+	}
+
+	performResult, err := stw.subtaskCtx.Provider.PerformAgentChain(ctx, taskID, subtaskID, msgChainID)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			ctx = context.Background()
+		}
+		errChainConsistency := stw.subtaskCtx.Provider.EnsureChainConsistency(ctx, msgChainID)
+		if errChainConsistency != nil {
+			err = errors.Join(err, errChainConsistency)
+		}
+		_ = stw.SetStatus(ctx, database.SubtaskStatusWaiting)
+		return fmt.Errorf("failed to perform agent chain for subtask %d: %w", subtaskID, err)
+	}
+
+	switch performResult {
+	case providers.PerformResultWaiting:
+		if err := stw.SetStatus(ctx, database.SubtaskStatusWaiting); err != nil {
+			stw.handleInterrupting(err)
+			return err
+		}
+	case providers.PerformResultDone:
+		if err := stw.SetStatus(ctx, database.SubtaskStatusFinished); err != nil {
+			stw.handleInterrupting(err)
+			return fmt.Errorf("failed to set subtask %d status to finished: %w", subtaskID, err)
+		}
+	case providers.PerformResultError:
+		if err := stw.SetStatus(ctx, database.SubtaskStatusFailed); err != nil {
+			stw.handleInterrupting(err)
+			return fmt.Errorf("failed to set subtask %d status to failed: %w", subtaskID, err)
+		}
+	default:
+		return fmt.Errorf("unknown perform result: %d", performResult)
+	}
+
+	return nil
+}
+
+// handleInterrupting sets this subtask (and task/flow via SetStatus back-propagation)
+// to Waiting when err is context.Canceled or context.DeadlineExceeded. Use after the subtask
+// was advanced past Waiting (e.g. Running) but the run aborts before PerformAgentChain's
+// normal error handler, or when a late SetStatus fails with a context interruption.
+func (stw *subtaskWorker) handleInterrupting(err error) {
+	if err == nil {
+		return
+	}
+	if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
+		return
+	}
+	if stw.IsCompleted() {
+		return
+	}
+
+	resetCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	if errSt := stw.SetStatus(resetCtx, database.SubtaskStatusWaiting); errSt != nil {
+		logrus.WithError(errSt).Warn("failed to set subtask waiting after run interrupt")
+	}
+}
+
+func (stw *subtaskWorker) Finish(ctx context.Context) error {
+	if stw.IsCompleted() {
+		return fmt.Errorf("subtask has already completed")
+	}
+
+	if err := stw.SetStatus(ctx, database.SubtaskStatusFinished); err != nil {
+		return err
+	}
+
+	return nil
+}

backend/pkg/controller/subtasks.go

@@ -0,0 +1,193 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"sync"
+
+	"pentagi/pkg/database"
+)
+
+type NewSubtaskInfo struct {
+	Title       string
+	Description string
+}
+
+type SubtaskController interface {
+	LoadSubtasks(ctx context.Context, taskID int64, updater TaskUpdater) error
+	GenerateSubtasks(ctx context.Context) error
+	RefineSubtasks(ctx context.Context) error
+	PopSubtask(ctx context.Context, updater TaskUpdater) (SubtaskWorker, error)
+	ListSubtasks(ctx context.Context) []SubtaskWorker
+	GetSubtask(ctx context.Context, subtaskID int64) (SubtaskWorker, error)
+}
+
+type subtaskController struct {
+	mx       *sync.Mutex
+	taskCtx  *TaskContext
+	subtasks map[int64]SubtaskWorker
+}
+
+func NewSubtaskController(taskCtx *TaskContext) SubtaskController {
+	return &subtaskController{
+		mx:       &sync.Mutex{},
+		taskCtx:  taskCtx,
+		subtasks: make(map[int64]SubtaskWorker),
+	}
+}
+
+func (stc *subtaskController) LoadSubtasks(ctx context.Context, taskID int64, updater TaskUpdater) error {
+	stc.mx.Lock()
+	defer stc.mx.Unlock()
+
+	subtasks, err := stc.taskCtx.DB.GetTaskSubtasks(ctx, taskID)
+	if err != nil {
+		return fmt.Errorf("failed to get subtasks for task %d: %w", taskID, err)
+	}
+
+	if len(subtasks) == 0 {
+		return fmt.Errorf("no subtasks found for task %d: %w", taskID, ErrNothingToLoad)
+	}
+
+	for _, subtask := range subtasks {
+		st, err := LoadSubtaskWorker(ctx, subtask, stc.taskCtx, updater)
+		if err != nil {
+			if errors.Is(err, ErrNothingToLoad) {
+				continue
+			}
+
+			return fmt.Errorf("failed to create subtask worker: %w", err)
+		}
+
+		stc.subtasks[subtask.ID] = st
+	}
+
+	return nil
+}
+
+func (stc *subtaskController) GenerateSubtasks(ctx context.Context) error {
+	plan, err := stc.taskCtx.Provider.GenerateSubtasks(ctx, stc.taskCtx.TaskID)
+	if err != nil {
+		return fmt.Errorf("failed to generate subtasks for task %d: %w", stc.taskCtx.TaskID, err)
+	}
+
+	if len(plan) == 0 {
+		return fmt.Errorf("no subtasks generated for task %d", stc.taskCtx.TaskID)
+	}
+
+	// TODO: change it to insert subtasks in transaction
+	for _, info := range plan {
+		_, err := stc.taskCtx.DB.CreateSubtask(ctx, 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
+}
+
+func (stc *subtaskController) RefineSubtasks(ctx context.Context) error {
+	subtasks, err := stc.taskCtx.DB.GetTaskSubtasks(ctx, stc.taskCtx.TaskID)
+	if err != nil {
+		return fmt.Errorf("failed to get task %d subtasks: %w", stc.taskCtx.TaskID, err)
+	}
+
+	plan, err := stc.taskCtx.Provider.RefineSubtasks(ctx, stc.taskCtx.TaskID)
+	if err != nil {
+		return fmt.Errorf("failed to refine subtasks for task %d: %w", stc.taskCtx.TaskID, err)
+	}
+
+	if len(plan) == 0 {
+		return nil // no subtasks refined
+	}
+
+	subtaskIDs := make([]int64, 0, len(subtasks))
+	for _, subtask := range subtasks {
+		if subtask.Status == database.SubtaskStatusCreated {
+			subtaskIDs = append(subtaskIDs, subtask.ID)
+		}
+	}
+
+	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)
+	}
+
+	// TODO: change it to insert subtasks in transaction and union it with delete ones
+	for _, info := range plan {
+		_, err := stc.taskCtx.DB.CreateSubtask(ctx, 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
+}
+
+func (stc *subtaskController) PopSubtask(ctx context.Context, updater TaskUpdater) (SubtaskWorker, error) {
+	stc.mx.Lock()
+	defer stc.mx.Unlock()
+
+	subtasks, err := stc.taskCtx.DB.GetTaskPlannedSubtasks(ctx, stc.taskCtx.TaskID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get task planned subtasks: %w", err)
+	}
+
+	if len(subtasks) == 0 {
+		return nil, nil
+	}
+
+	stdb := subtasks[0]
+	if st, ok := stc.subtasks[stdb.ID]; ok {
+		return st, nil
+	}
+
+	st, err := NewSubtaskWorker(ctx, stc.taskCtx, stdb.ID, stdb.Title, stdb.Description, updater)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create subtask worker: %w", err)
+	}
+
+	stc.subtasks[stdb.ID] = st
+
+	return st, nil
+}
+
+func (stc *subtaskController) ListSubtasks(ctx context.Context) []SubtaskWorker {
+	stc.mx.Lock()
+	defer stc.mx.Unlock()
+
+	subtasks := make([]SubtaskWorker, 0)
+	for _, subtask := range stc.subtasks {
+		subtasks = append(subtasks, subtask)
+	}
+
+	sort.Slice(subtasks, func(i, j int) bool {
+		return subtasks[i].GetSubtaskID() < subtasks[j].GetSubtaskID()
+	})
+
+	return subtasks
+}
+
+func (stc *subtaskController) GetSubtask(ctx context.Context, subtaskID int64) (SubtaskWorker, error) {
+	stc.mx.Lock()
+	defer stc.mx.Unlock()
+
+	subtask, ok := stc.subtasks[subtaskID]
+	if !ok {
+		return nil, fmt.Errorf("subtask not found")
+	}
+
+	return subtask, nil
+}

backend/pkg/controller/task.go

@@ -0,0 +1,399 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"pentagi/pkg/database"
+	obs "pentagi/pkg/observability"
+	"pentagi/pkg/providers"
+	"pentagi/pkg/tools"
+
+	"github.com/sirupsen/logrus"
+)
+
+type FlowUpdater interface {
+	SetStatus(ctx context.Context, status database.FlowStatus) error
+}
+
+type TaskWorker interface {
+	GetTaskID() int64
+	GetFlowID() int64
+	GetUserID() int64
+	GetTitle() string
+	IsCompleted() bool
+	IsWaiting() bool
+	GetStatus(ctx context.Context) (database.TaskStatus, error)
+	SetStatus(ctx context.Context, status database.TaskStatus) error
+	GetResult(ctx context.Context) (string, error)
+	SetResult(ctx context.Context, result string) error
+	PutInput(ctx context.Context, input string) error
+	Run(ctx context.Context) error
+	Finish(ctx context.Context) error
+}
+
+type taskWorker struct {
+	mx        *sync.RWMutex
+	stc       SubtaskController
+	taskCtx   *TaskContext
+	updater   FlowUpdater
+	completed bool
+	waiting   bool
+}
+
+func NewTaskWorker(
+	ctx context.Context,
+	flowCtx *FlowContext,
+	input string,
+	updater FlowUpdater,
+) (TaskWorker, error) {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.NewTaskWorker")
+	defer span.End()
+
+	ctx = tools.PutAgentContext(ctx, database.MsgchainTypePrimaryAgent)
+
+	title, err := flowCtx.Provider.GetTaskTitle(ctx, input)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get task title: %w", err)
+	}
+
+	task, err := flowCtx.DB.CreateTask(ctx, database.CreateTaskParams{
+		Status: database.TaskStatusCreated,
+		Title:  title,
+		Input:  input,
+		FlowID: flowCtx.FlowID,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to create task in DB: %w", err)
+	}
+
+	flowCtx.Publisher.TaskCreated(ctx, task, []database.Subtask{})
+
+	taskCtx := &TaskContext{
+		FlowContext: *flowCtx,
+		TaskID:      task.ID,
+		TaskTitle:   title,
+		TaskInput:   input,
+	}
+	stc := NewSubtaskController(taskCtx)
+
+	_, err = taskCtx.MsgLog.PutTaskMsg(
+		ctx,
+		database.MsglogTypeInput,
+		taskCtx.TaskID,
+		"", // thinking is empty because this is input
+		input,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to put input for task %d: %w", taskCtx.TaskID, err)
+	}
+
+	err = stc.GenerateSubtasks(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate subtasks: %w", err)
+	}
+
+	subtasks, err := flowCtx.DB.GetTaskSubtasks(ctx, task.ID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get subtasks for task %d: %w", task.ID, err)
+	}
+
+	flowCtx.Publisher.TaskUpdated(ctx, task, subtasks)
+
+	return &taskWorker{
+		mx:        &sync.RWMutex{},
+		stc:       stc,
+		taskCtx:   taskCtx,
+		updater:   updater,
+		completed: false,
+		waiting:   false,
+	}, nil
+}
+
+func LoadTaskWorker(
+	ctx context.Context,
+	task database.Task,
+	flowCtx *FlowContext,
+	updater FlowUpdater,
+) (TaskWorker, error) {
+	ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.LoadTaskWorker")
+	defer span.End()
+
+	ctx = tools.PutAgentContext(ctx, database.MsgchainTypePrimaryAgent)
+	taskCtx := &TaskContext{
+		FlowContext: *flowCtx,
+		TaskID:      task.ID,
+		TaskTitle:   task.Title,
+		TaskInput:   task.Input,
+	}
+
+	stc := NewSubtaskController(taskCtx)
+	var completed, waiting bool
+	switch task.Status {
+	case database.TaskStatusFinished, database.TaskStatusFailed:
+		completed = true
+	case database.TaskStatusWaiting:
+		waiting = true
+	case database.TaskStatusRunning:
+	case database.TaskStatusCreated:
+		return nil, fmt.Errorf("task %d has created yet: loading aborted: %w", task.ID, ErrNothingToLoad)
+	}
+
+	tw := &taskWorker{
+		mx:        &sync.RWMutex{},
+		stc:       stc,
+		taskCtx:   taskCtx,
+		updater:   updater,
+		completed: completed,
+		waiting:   waiting,
+	}
+
+	if err := tw.stc.LoadSubtasks(ctx, task.ID, tw); err != nil {
+		return nil, fmt.Errorf("failed to load subtasks for task %d: %w", task.ID, err)
+	}
+
+	return tw, nil
+}
+
+func (tw *taskWorker) GetTaskID() int64 {
+	return tw.taskCtx.TaskID
+}
+
+func (tw *taskWorker) GetFlowID() int64 {
+	return tw.taskCtx.FlowID
+}
+
+func (tw *taskWorker) GetUserID() int64 {
+	return tw.taskCtx.UserID
+}
+
+func (tw *taskWorker) GetTitle() string {
+	return tw.taskCtx.TaskTitle
+}
+
+func (tw *taskWorker) IsCompleted() bool {
+	tw.mx.RLock()
+	defer tw.mx.RUnlock()
+
+	return tw.completed
+}
+
+func (tw *taskWorker) IsWaiting() bool {
+	tw.mx.RLock()
+	defer tw.mx.RUnlock()
+
+	return tw.waiting
+}
+
+func (tw *taskWorker) GetStatus(ctx context.Context) (database.TaskStatus, error) {
+	task, err := tw.taskCtx.DB.GetTask(ctx, tw.taskCtx.TaskID)
+	if err != nil {
+		return database.TaskStatusFailed, err
+	}
+
+	return task.Status, nil
+}
+
+// this function is exclusively change task internal properties "completed" and "waiting"
+func (tw *taskWorker) SetStatus(ctx context.Context, status database.TaskStatus) error {
+	task, err := tw.taskCtx.DB.UpdateTaskStatus(ctx, database.UpdateTaskStatusParams{
+		Status: status,
+		ID:     tw.taskCtx.TaskID,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to set task %d status: %w", tw.taskCtx.TaskID, err)
+	}
+
+	subtasks, err := tw.taskCtx.DB.GetTaskSubtasks(ctx, tw.taskCtx.TaskID)
+	if err != nil {
+		return fmt.Errorf("failed to get task %d subtasks: %w", tw.taskCtx.TaskID, err)
+	}
+
+	tw.taskCtx.Publisher.TaskUpdated(ctx, task, subtasks)
+
+	tw.mx.Lock()
+	defer tw.mx.Unlock()
+
+	switch status {
+	case database.TaskStatusRunning:
+		tw.completed = false
+		tw.waiting = false
+		err = tw.updater.SetStatus(ctx, database.FlowStatusRunning)
+	case database.TaskStatusWaiting:
+		tw.completed = false
+		tw.waiting = true
+		err = tw.updater.SetStatus(ctx, database.FlowStatusWaiting)
+	case database.TaskStatusFinished, database.TaskStatusFailed:
+		tw.completed = true
+		tw.waiting = false
+		// the last task was done, set flow status to Waiting new user input
+		err = tw.updater.SetStatus(ctx, database.FlowStatusWaiting)
+	default:
+		// status Created is not possible to set by this call
+		return fmt.Errorf("unsupported task status: %s", status)
+	}
+	if err != nil {
+		return fmt.Errorf("failed to set flow status in back propagation: %w", err)
+	}
+
+	return nil
+}
+
+func (tw *taskWorker) GetResult(ctx context.Context) (string, error) {
+	task, err := tw.taskCtx.DB.GetTask(ctx, tw.taskCtx.TaskID)
+	if err != nil {
+		return "", err
+	}
+
+	return task.Result, nil
+}
+
+func (tw *taskWorker) SetResult(ctx context.Context, result string) error {
+	_, err := tw.taskCtx.DB.UpdateTaskResult(ctx, database.UpdateTaskResultParams{
+		Result: result,
+		ID:     tw.taskCtx.TaskID,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to set task %d result: %w", tw.taskCtx.TaskID, err)
+	}
+
+	return nil
+}
+
+func (tw *taskWorker) PutInput(ctx context.Context, input string) error {
+	if !tw.IsWaiting() {
+		return fmt.Errorf("task is not waiting")
+	}
+
+	for _, st := range tw.stc.ListSubtasks(ctx) {
+		if !st.IsCompleted() && st.IsWaiting() {
+			if err := st.PutInput(ctx, input); err != nil {
+				return fmt.Errorf("failed to put input to subtask %d: %w", st.GetSubtaskID(), err)
+			} else {
+				break
+			}
+		}
+	}
+
+	return nil
+}
+
+func (tw *taskWorker) Run(ctx context.Context) error {
+	ctx = tools.PutAgentContext(ctx, database.MsgchainTypePrimaryAgent)
+
+	for len(tw.stc.ListSubtasks(ctx)) < providers.TasksNumberLimit+3 {
+		st, err := tw.stc.PopSubtask(ctx, tw)
+		if err != nil {
+			tw.handleInterrupting(err)
+			return err
+		}
+
+		// empty queue for subtasks means that task is done
+		if st == nil {
+			break
+		}
+
+		if err := st.Run(ctx); err != nil {
+			tw.handleInterrupting(err)
+			return err
+		}
+
+		// pass through if task is waiting from back status propagation
+		if tw.IsWaiting() {
+			return nil
+		} // otherwise subtask is done
+
+		if err := tw.stc.RefineSubtasks(ctx); err != nil {
+			if errors.Is(err, context.Canceled) {
+				ctx = context.Background()
+			}
+			_ = tw.SetStatus(ctx, database.TaskStatusWaiting)
+			return fmt.Errorf("failed to refine subtasks list for the task %d: %w", tw.taskCtx.TaskID, err)
+		}
+	}
+
+	jobResult, err := tw.taskCtx.Provider.GetTaskResult(ctx, tw.taskCtx.TaskID)
+	if err != nil {
+		tw.handleInterrupting(err)
+		return fmt.Errorf("failed to get task %d result: %w", tw.taskCtx.TaskID, err)
+	}
+
+	var taskStatus database.TaskStatus
+	if jobResult.Success {
+		taskStatus = database.TaskStatusFinished
+	} else {
+		taskStatus = database.TaskStatusFailed
+	}
+
+	if err := tw.SetResult(ctx, jobResult.Result); err != nil {
+		tw.handleInterrupting(err)
+		return err
+	}
+
+	if err := tw.SetStatus(ctx, taskStatus); err != nil {
+		tw.handleInterrupting(err)
+		return err
+	}
+
+	format := database.MsglogResultFormatMarkdown
+	_, err = tw.taskCtx.MsgLog.PutTaskMsgResult(
+		ctx,
+		database.MsglogTypeReport,
+		tw.taskCtx.TaskID,
+		"", // thinking is empty because agent can't return it
+		tw.taskCtx.TaskTitle,
+		jobResult.Result,
+		format,
+	)
+	if err != nil {
+		tw.handleInterrupting(err)
+		return fmt.Errorf("failed to put report for task %d: %w", tw.taskCtx.TaskID, err)
+	}
+
+	return nil
+}
+
+// handleInterrupting sets this task (and the flow via taskWorker.SetStatus)
+// to Waiting when err is context.Canceled or context.DeadlineExceeded. Skips if the task is
+// already marked completed in memory (Finished/Failed) so we do not revive a finished task.
+func (tw *taskWorker) handleInterrupting(err error) {
+	if err == nil {
+		return
+	}
+	if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
+		return
+	}
+	if tw.IsCompleted() {
+		return
+	}
+
+	resetCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	if errSt := tw.SetStatus(resetCtx, database.TaskStatusWaiting); errSt != nil {
+		logrus.WithError(errSt).Warn("failed to set task waiting after run interrupt")
+	}
+}
+
+func (tw *taskWorker) Finish(ctx context.Context) error {
+	if tw.IsCompleted() {
+		return fmt.Errorf("task has already completed")
+	}
+
+	for _, st := range tw.stc.ListSubtasks(ctx) {
+		if !st.IsCompleted() {
+			if err := st.Finish(ctx); err != nil {
+				return err
+			}
+		}
+	}
+
+	if err := tw.SetStatus(ctx, database.TaskStatusFinished); err != nil {
+		return err
+	}
+
+	return nil
+}

backend/pkg/controller/tasks.go

@@ -0,0 +1,110 @@
+package controller
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"sync"
+)
+
+type TaskController interface {
+	CreateTask(ctx context.Context, input string, updater FlowUpdater) (TaskWorker, error)
+	LoadTasks(ctx context.Context, flowID int64, updater FlowUpdater) error
+	ListTasks(ctx context.Context) []TaskWorker
+	GetTask(ctx context.Context, taskID int64) (TaskWorker, error)
+}
+
+type taskController struct {
+	mx      *sync.Mutex
+	tasks   map[int64]TaskWorker
+	updater FlowUpdater
+	flowCtx *FlowContext
+}
+
+func NewTaskController(flowCtx *FlowContext) TaskController {
+	return &taskController{
+		mx:      &sync.Mutex{},
+		tasks:   make(map[int64]TaskWorker),
+		flowCtx: flowCtx,
+	}
+}
+
+func (tc *taskController) LoadTasks(
+	ctx context.Context,
+	flowID int64,
+	updater FlowUpdater,
+) error {
+	tc.mx.Lock()
+	defer tc.mx.Unlock()
+
+	tasks, err := tc.flowCtx.DB.GetFlowTasks(ctx, flowID)
+	if err != nil {
+		return fmt.Errorf("failed to get flow tasks: %w", err)
+	}
+
+	if len(tasks) == 0 {
+		return fmt.Errorf("no tasks found for flow %d: %w", flowID, ErrNothingToLoad)
+	}
+
+	for _, task := range tasks {
+		tw, err := LoadTaskWorker(ctx, task, tc.flowCtx, updater)
+		if err != nil {
+			if errors.Is(err, ErrNothingToLoad) {
+				continue
+			}
+
+			return fmt.Errorf("failed to load task worker: %w", err)
+		}
+
+		tc.tasks[task.ID] = tw
+	}
+
+	return nil
+}
+
+func (tc *taskController) CreateTask(
+	ctx context.Context,
+	input string,
+	updater FlowUpdater,
+) (TaskWorker, error) {
+	tc.mx.Lock()
+	defer tc.mx.Unlock()
+
+	tw, err := NewTaskWorker(ctx, tc.flowCtx, input, updater)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create task worker: %w", err)
+	}
+
+	tc.tasks[tw.GetTaskID()] = tw
+
+	return tw, nil
+}
+
+func (tc *taskController) ListTasks(ctx context.Context) []TaskWorker {
+	tc.mx.Lock()
+	defer tc.mx.Unlock()
+
+	tasks := make([]TaskWorker, 0)
+	for _, task := range tc.tasks {
+		tasks = append(tasks, task)
+	}
+
+	sort.Slice(tasks, func(i, j int) bool {
+		return tasks[i].GetTaskID() < tasks[j].GetTaskID()
+	})
+
+	return tasks
+}
+
+func (tc *taskController) GetTask(ctx context.Context, taskID int64) (TaskWorker, error) {
+	tc.mx.Lock()
+	defer tc.mx.Unlock()
+
+	task, ok := tc.tasks[taskID]
+	if !ok {
+		return nil, fmt.Errorf("task %d not found", taskID)
+	}
+
+	return task, nil
+}

backend/pkg/controller/tclog.go

@@ -0,0 +1,126 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type FlowToolCallLogWorker interface {
+	PutLog(
+		ctx context.Context,
+		callID string,
+		name string,
+		args json.RawMessage,
+		taskID *int64,
+		subtaskID *int64,
+	) (int64, error)
+	UpdateLogSuccess(ctx context.Context, id int64, result string, durationSeconds float64) error
+	UpdateLogFailed(ctx context.Context, id int64, result string, durationSeconds float64) error
+	GetLog(ctx context.Context, id int64) (database.Toolcall, error)
+}
+
+type flowToolCallLogWorker struct {
+	db     database.Querier
+	mx     *sync.Mutex
+	flowID int64
+	pub    subscriptions.FlowPublisher
+}
+
+func NewFlowToolCallLogWorker(
+	db database.Querier,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) FlowToolCallLogWorker {
+	return &flowToolCallLogWorker{
+		db:     db,
+		mx:     &sync.Mutex{},
+		flowID: flowID,
+		pub:    pub,
+	}
+}
+
+func (w *flowToolCallLogWorker) PutLog(
+	ctx context.Context,
+	callID string,
+	name string,
+	args json.RawMessage,
+	taskID *int64,
+	subtaskID *int64,
+) (int64, error) {
+	w.mx.Lock()
+	defer w.mx.Unlock()
+
+	tc, err := w.db.CreateToolcall(ctx, database.CreateToolcallParams{
+		CallID:    callID,
+		Status:    database.ToolcallStatusRunning,
+		Name:      name,
+		Args:      args,
+		FlowID:    w.flowID,
+		TaskID:    database.Int64ToNullInt64(taskID),
+		SubtaskID: database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to create tool call log: %w", err)
+	}
+
+	w.pub.ToolCallLogAdded(ctx, tc)
+
+	return tc.ID, nil
+}
+
+func (w *flowToolCallLogWorker) UpdateLogSuccess(
+	ctx context.Context,
+	id int64,
+	result string,
+	durationSeconds float64,
+) error {
+	tc, err := w.db.UpdateToolcallFinishedResult(ctx, database.UpdateToolcallFinishedResultParams{
+		Result:          result,
+		DurationSeconds: durationSeconds,
+		ID:              id,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to update tool call log result: %w", err)
+	}
+
+	w.pub.ToolCallLogUpdated(ctx, tc)
+
+	return nil
+}
+
+func (w *flowToolCallLogWorker) UpdateLogFailed(
+	ctx context.Context,
+	id int64,
+	result string,
+	durationSeconds float64,
+) error {
+	tc, err := w.db.UpdateToolcallFailedResult(ctx, database.UpdateToolcallFailedResultParams{
+		Result:          result,
+		DurationSeconds: durationSeconds,
+		ID:              id,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to update tool call log failed result: %w", err)
+	}
+
+	w.pub.ToolCallLogUpdated(ctx, tc)
+
+	return nil
+}
+
+func (w *flowToolCallLogWorker) GetLog(ctx context.Context, id int64) (database.Toolcall, error) {
+	tc, err := w.db.GetFlowToolcall(ctx, database.GetFlowToolcallParams{
+		ID:     id,
+		FlowID: w.flowID,
+	})
+	if err != nil {
+		return database.Toolcall{}, fmt.Errorf("failed to get tool call log: %w", err)
+	}
+
+	return tc, nil
+}

backend/pkg/controller/tclogs.go

@@ -0,0 +1,71 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type ToolCallLogController interface {
+	NewFlowToolCallLog(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowToolCallLogWorker, error)
+	ListFlowsToolCallLog(ctx context.Context) ([]FlowToolCallLogWorker, error)
+	GetFlowToolCallLog(ctx context.Context, flowID int64) (FlowToolCallLogWorker, error)
+}
+
+type toolCallLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowToolCallLogWorker
+}
+
+func NewToolCallLogController(db database.Querier) ToolCallLogController {
+	return &toolCallLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowToolCallLogWorker),
+	}
+}
+
+func (c *toolCallLogController) NewFlowToolCallLog(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowToolCallLogWorker, error) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	flw := NewFlowToolCallLogWorker(c.db, flowID, pub)
+	c.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (c *toolCallLogController) ListFlowsToolCallLog(ctx context.Context) ([]FlowToolCallLogWorker, error) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	flows := make([]FlowToolCallLogWorker, 0, len(c.flows))
+	for _, flw := range c.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (c *toolCallLogController) GetFlowToolCallLog(
+	ctx context.Context,
+	flowID int64,
+) (FlowToolCallLogWorker, error) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	flw, ok := c.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}

backend/pkg/controller/termlog.go

@@ -0,0 +1,100 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type FlowTermLogWorker interface {
+	PutMsg(
+		ctx context.Context,
+		msgType database.TermlogType,
+		msg string,
+		containerID int64,
+		taskID, subtaskID *int64,
+	) (int64, error)
+	GetMsg(ctx context.Context, msgID int64) (database.Termlog, error)
+	GetContainers(ctx context.Context) ([]database.Container, error)
+}
+
+type flowTermLogWorker struct {
+	db         database.Querier
+	mx         *sync.Mutex
+	flowID     int64
+	containers map[int64]struct{}
+	pub        subscriptions.FlowPublisher
+}
+
+func NewFlowTermLogWorker(db database.Querier, flowID int64, pub subscriptions.FlowPublisher) FlowTermLogWorker {
+	return &flowTermLogWorker{
+		db:         db,
+		mx:         &sync.Mutex{},
+		flowID:     flowID,
+		containers: make(map[int64]struct{}),
+		pub:        pub,
+	}
+}
+
+func (tlw *flowTermLogWorker) PutMsg(
+	ctx context.Context,
+	msgType database.TermlogType,
+	msg string,
+	containerID int64,
+	taskID, subtaskID *int64,
+) (int64, error) {
+	tlw.mx.Lock()
+	defer tlw.mx.Unlock()
+
+	if _, ok := tlw.containers[containerID]; !ok {
+		// try to update the container map
+		containers, err := tlw.GetContainers(ctx)
+		if err != nil {
+			return 0, err
+		}
+		tlw.containers = make(map[int64]struct{})
+		for _, container := range containers {
+			tlw.containers[container.ID] = struct{}{}
+		}
+		if _, ok := tlw.containers[containerID]; !ok {
+			return 0, fmt.Errorf("container not found")
+		}
+	}
+
+	termLog, err := tlw.db.CreateTermLog(ctx, database.CreateTermLogParams{
+		Type:        msgType,
+		Text:        database.SanitizeUTF8(msg),
+		ContainerID: containerID,
+		FlowID:      tlw.flowID,
+		TaskID:      database.Int64ToNullInt64(taskID),
+		SubtaskID:   database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to create termlog: %w", err)
+	}
+
+	tlw.pub.TerminalLogAdded(ctx, termLog)
+
+	return termLog.ID, nil
+}
+
+func (tlw *flowTermLogWorker) GetMsg(ctx context.Context, msgID int64) (database.Termlog, error) {
+	msg, err := tlw.db.GetTermLog(ctx, msgID)
+	if err != nil {
+		return database.Termlog{}, fmt.Errorf("failed to get termlog: %w", err)
+	}
+
+	return msg, nil
+}
+
+func (tlw *flowTermLogWorker) GetContainers(ctx context.Context) ([]database.Container, error) {
+	containers, err := tlw.db.GetFlowContainers(ctx, tlw.flowID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get containers: %w", err)
+	}
+
+	return containers, nil
+}

backend/pkg/controller/termlogs.go

@@ -0,0 +1,81 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type TermLogController interface {
+	NewFlowTermLog(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowTermLogWorker, error)
+	ListFlowsTermLog(ctx context.Context) ([]FlowTermLogWorker, error)
+	GetFlowTermLog(ctx context.Context, flowID int64) (FlowTermLogWorker, error)
+	GetFlowContainers(ctx context.Context, flowID int64) ([]database.Container, error)
+}
+
+type termLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowTermLogWorker
+}
+
+func NewTermLogController(db database.Querier) TermLogController {
+	return &termLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowTermLogWorker),
+	}
+}
+
+func (tlc *termLogController) NewFlowTermLog(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowTermLogWorker, error) {
+	tlc.mx.Lock()
+	defer tlc.mx.Unlock()
+
+	flw := NewFlowTermLogWorker(tlc.db, flowID, pub)
+	tlc.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (tlc *termLogController) ListFlowsTermLog(ctx context.Context) ([]FlowTermLogWorker, error) {
+	tlc.mx.Lock()
+	defer tlc.mx.Unlock()
+
+	flows := make([]FlowTermLogWorker, 0, len(tlc.flows))
+	for _, flw := range tlc.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (tlc *termLogController) GetFlowTermLog(ctx context.Context, flowID int64) (FlowTermLogWorker, error) {
+	tlc.mx.Lock()
+	defer tlc.mx.Unlock()
+
+	flw, ok := tlc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}
+
+func (tlc *termLogController) GetFlowContainers(ctx context.Context, flowID int64) ([]database.Container, error) {
+	tlc.mx.Lock()
+	defer tlc.mx.Unlock()
+
+	flw, ok := tlc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw.GetContainers(ctx)
+}

backend/pkg/controller/vslog.go

@@ -0,0 +1,93 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type FlowVectorStoreLogWorker interface {
+	PutLog(
+		ctx context.Context,
+		initiator database.MsgchainType,
+		executor database.MsgchainType,
+		filter json.RawMessage,
+		query string,
+		action database.VecstoreActionType,
+		result string,
+		taskID *int64,
+		subtaskID *int64,
+	) (int64, error)
+	GetLog(ctx context.Context, msgID int64) (database.Vecstorelog, error)
+}
+
+type flowVectorStoreLogWorker struct {
+	db         database.Querier
+	mx         *sync.Mutex
+	flowID     int64
+	containers map[int64]struct{}
+	pub        subscriptions.FlowPublisher
+}
+
+func NewFlowVectorStoreLogWorker(
+	db database.Querier,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) FlowVectorStoreLogWorker {
+	return &flowVectorStoreLogWorker{
+		db:     db,
+		mx:     &sync.Mutex{},
+		flowID: flowID,
+		pub:    pub,
+	}
+}
+
+func (vslw *flowVectorStoreLogWorker) PutLog(
+	ctx context.Context,
+	initiator database.MsgchainType,
+	executor database.MsgchainType,
+	filter json.RawMessage,
+	query string,
+	action database.VecstoreActionType,
+	result string,
+	taskID *int64,
+	subtaskID *int64,
+) (int64, error) {
+	vslw.mx.Lock()
+	defer vslw.mx.Unlock()
+
+	vsLog, err := vslw.db.CreateVectorStoreLog(ctx, database.CreateVectorStoreLogParams{
+		Initiator: initiator,
+		Executor:  executor,
+		Filter:    filter,
+		Query:     query,
+		Action:    action,
+		Result:    result,
+		FlowID:    vslw.flowID,
+		TaskID:    database.Int64ToNullInt64(taskID),
+		SubtaskID: database.Int64ToNullInt64(subtaskID),
+	})
+	if err != nil {
+		return 0, fmt.Errorf("failed to create vector store log: %w", err)
+	}
+
+	vslw.pub.VectorStoreLogAdded(ctx, vsLog)
+
+	return vsLog.ID, nil
+}
+
+func (vslw *flowVectorStoreLogWorker) GetLog(ctx context.Context, msgID int64) (database.Vecstorelog, error) {
+	msg, err := vslw.db.GetFlowVectorStoreLog(ctx, database.GetFlowVectorStoreLogParams{
+		ID:     msgID,
+		FlowID: vslw.flowID,
+	})
+	if err != nil {
+		return database.Vecstorelog{}, fmt.Errorf("failed to get vector store log: %w", err)
+	}
+
+	return msg, nil
+}

backend/pkg/controller/vslogs.go

@@ -0,0 +1,71 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"pentagi/pkg/database"
+	"pentagi/pkg/graph/subscriptions"
+)
+
+type VectorStoreLogController interface {
+	NewFlowVectorStoreLog(ctx context.Context, flowID int64, pub subscriptions.FlowPublisher) (FlowVectorStoreLogWorker, error)
+	ListFlowsVectorStoreLog(ctx context.Context) ([]FlowVectorStoreLogWorker, error)
+	GetFlowVectorStoreLog(ctx context.Context, flowID int64) (FlowVectorStoreLogWorker, error)
+}
+
+type vectorStoreLogController struct {
+	db    database.Querier
+	mx    *sync.Mutex
+	flows map[int64]FlowVectorStoreLogWorker
+}
+
+func NewVectorStoreLogController(db database.Querier) VectorStoreLogController {
+	return &vectorStoreLogController{
+		db:    db,
+		mx:    &sync.Mutex{},
+		flows: make(map[int64]FlowVectorStoreLogWorker),
+	}
+}
+
+func (vslc *vectorStoreLogController) NewFlowVectorStoreLog(
+	ctx context.Context,
+	flowID int64,
+	pub subscriptions.FlowPublisher,
+) (FlowVectorStoreLogWorker, error) {
+	vslc.mx.Lock()
+	defer vslc.mx.Unlock()
+
+	flw := NewFlowVectorStoreLogWorker(vslc.db, flowID, pub)
+	vslc.flows[flowID] = flw
+
+	return flw, nil
+}
+
+func (tlc *vectorStoreLogController) ListFlowsVectorStoreLog(ctx context.Context) ([]FlowVectorStoreLogWorker, error) {
+	tlc.mx.Lock()
+	defer tlc.mx.Unlock()
+
+	flows := make([]FlowVectorStoreLogWorker, 0, len(tlc.flows))
+	for _, flw := range tlc.flows {
+		flows = append(flows, flw)
+	}
+
+	return flows, nil
+}
+
+func (vslc *vectorStoreLogController) GetFlowVectorStoreLog(
+	ctx context.Context,
+	flowID int64,
+) (FlowVectorStoreLogWorker, error) {
+	vslc.mx.Lock()
+	defer vslc.mx.Unlock()
+
+	flw, ok := vslc.flows[flowID]
+	if !ok {
+		return nil, fmt.Errorf("flow not found")
+	}
+
+	return flw, nil
+}