feat(flows): implement provider renaming and reset functionality

- Added `RenameFlowsProvider` and `ResetFlowsProviderToDefault` methods to the `FlowController` to handle renaming of user-defined providers and resetting to built-in defaults.
- Introduced SQL queries for bulk updating flow and assistant provider names based on user actions, ensuring idempotency and safe retries.
- Enhanced error handling and logging for provider updates, ensuring that flows and assistants remain valid after provider changes.
- Added unit tests to verify the correct behavior of provider renaming and resetting functionalities.
This commit is contained in:
Dmitry Ng
2026-08-04 01:06:58 +03:00
parent 2da929d38d
commit db68cde21c
15 changed files with 849 additions and 33 deletions
+41 -14
View File
@@ -482,7 +482,15 @@ func LoadFlowWorker(ctx context.Context, flow database.Flow, fwc flowWorkerCtx)
if errors.Is(err, ErrNothingToLoad) {
continue
}
return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to load assistant worker", err)
// One unloadable assistant must not take its flow down with it.
// Aborting here would leave the flow absent from the controller's
// map while its row stays alive in the DB, which makes the whole
// flow permanently unreachable ("flow not found" on every action) —
// the exact failure this used to produce when an assistant pointed
// at a renamed or deleted provider. The assistant simply stays
// unloaded and returns on the next start once its cause is fixed.
logger.WithError(err).Errorf("failed to load assistant %d, skipping it", assistant.ID)
continue
}
if err := fw.AddAssistant(ctx, aw); err != nil {
return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to add assistant worker", err)
@@ -860,7 +868,24 @@ func (fw *flowWorker) Rename(ctx context.Context, title string) error {
return nil
}
// switchProvider performs runtime provider switch for the flow
// switchProvider performs runtime provider switch for the flow.
//
// This is the single place where a running flow picks up a provider change. A
// rename or deletion of a user provider only rewrites the DB reference (see
// flowController.reassignFlowsProvider); the in-memory instance is refreshed
// here, on the next user input, or rebuilt from the DB row on the next start.
//
// Deciding whether anything changed is delegated to SetProvider, which compares
// the raw configuration and not just the provider name — a user provider may be
// named exactly like a built-in one, so the name alone cannot tell an override
// apart from the default it shadows.
//
// Note on tool_call_id_template: it is resolved once, for a single model of the
// provider configuration. A provider that routes different models to different
// upstream backends (an OpenRouter-style gateway) may therefore need different
// templates per agent, and this single value can be wrong for some of them.
// Fixing that properly means keeping a per-model template registry, which is out
// of scope here; if real users hit it, this is the place to start.
func (fw *flowWorker) switchProvider(ctx context.Context, prv provider.Provider) error {
ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "controller.flowWorker.switchProvider")
defer span.End()
@@ -870,29 +895,31 @@ func (fw *flowWorker) switchProvider(ctx context.Context, prv provider.Provider)
}
logger := fw.logger.WithFields(logrus.Fields{
"old_provider_name": fw.flowCtx.Provider.Name().String(),
"old_provider_type": fw.flowCtx.Provider.Type().String(),
"new_provider_name": prv.Name().String(),
"new_provider_type": prv.Type().String(),
})
if fw.flowCtx.Provider.Name() == prv.Name() {
changed, tcIDTemplate, err := fw.flowCtx.Provider.SetProvider(ctx, prv)
if err != nil {
logger.WithError(err).Error("failed to set provider")
return fmt.Errorf("failed to set provider: %w", err)
}
if !changed {
logger.Debug("provider is the same, skipping switch")
return nil
}
logger.Info("switching flow provider")
if err := fw.flowCtx.Provider.SetProvider(ctx, prv); err != nil {
logger.WithError(err).Error("failed to set provider")
return fmt.Errorf("failed to set provider: %w", err)
}
// Every persisted value is taken from prv (and the template SetProvider
// resolved for it) rather than re-read from the shared flow provider, so a
// concurrent switch cannot interleave into a mixed-provider row.
flow, err := fw.flowCtx.DB.UpdateFlowProvider(ctx, database.UpdateFlowProviderParams{
ModelProviderName: prv.Name().String(),
ModelProviderType: database.ProviderType(prv.Type()),
ToolCallIDTemplate: fw.flowCtx.Provider.ToolCallIDTemplate(),
Model: fw.flowCtx.Provider.Model(pconfig.OptionsTypePrimaryAgent),
ToolCallIDTemplate: tcIDTemplate,
Model: prv.Model(pconfig.OptionsTypePrimaryAgent),
ID: fw.flowCtx.FlowID,
})
if err != nil {
@@ -901,8 +928,8 @@ func (fw *flowWorker) switchProvider(ctx context.Context, prv provider.Provider)
}
logger.WithFields(logrus.Fields{
"new_tool_call_id_template": fw.flowCtx.Provider.ToolCallIDTemplate(),
"new_model": fw.flowCtx.Provider.Model(pconfig.OptionsTypePrimaryAgent),
"new_tool_call_id_template": tcIDTemplate,
"new_model": prv.Model(pconfig.OptionsTypePrimaryAgent),
}).Info("provider switched successfully")
if containers, err := fw.flowCtx.DB.GetFlowContainers(ctx, fw.flowCtx.FlowID); err == nil {
+141
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"sort"
"sync"
"time"
"pentagi/pkg/config"
"pentagi/pkg/database"
@@ -50,8 +51,20 @@ type FlowController interface {
StopFlow(ctx context.Context, flowID int64) error
FinishFlow(ctx context.Context, flowID int64) error
RenameFlow(ctx context.Context, flowID int64, title string) error
RenameFlowsProvider(ctx context.Context, userID int64, oldName, newName provider.ProviderName) error
ResetFlowsProviderToDefault(
ctx context.Context,
userID int64,
oldName provider.ProviderName,
prvtype provider.ProviderType,
) error
}
// reassignProviderTimeout bounds the provider reference sweep. It is generous
// for two indexed UPDATEs and only exists so a stuck database cannot pin the
// goroutine forever once the sweep is detached from the request context.
const reassignProviderTimeout = 30 * time.Second
type flowController struct {
db database.Querier
mx *sync.Mutex
@@ -389,3 +402,131 @@ func (fc *flowController) RenameFlow(ctx context.Context, flowID int64, title st
return flow.Rename(ctx, title)
}
// RenameFlowsProvider repoints every flow and assistant of userID that still
// refers to oldName at newName, after the user renamed a custom LLM provider.
func (fc *flowController) RenameFlowsProvider(
ctx context.Context,
userID int64,
oldName, newName provider.ProviderName,
) error {
return fc.reassignFlowsProvider(ctx, userID, oldName, newName)
}
// ResetFlowsProviderToDefault repoints every flow and assistant of userID that
// referred to a just-deleted custom LLM provider at the built-in name for its
// type, which is literally the type string ("qwen", "openai", ...) — see
// provider.DefaultProviderName*. That name always resolves, so the flow stays
// loadable instead of failing with "provider not found by name".
func (fc *flowController) ResetFlowsProviderToDefault(
ctx context.Context,
userID int64,
oldName provider.ProviderName,
prvtype provider.ProviderType,
) error {
return fc.reassignFlowsProvider(ctx, userID, oldName, provider.ProviderName(prvtype))
}
// reassignFlowsProvider rewrites the provider reference stored on a user's flow
// and assistant rows. It deliberately does *not* touch loaded workers:
//
// - Nothing here blocks on an LLM. Building a provider instance probes the
// upstream API to resolve a tool call ID template, so switching loaded
// workers inline would tie a "rename provider" click to LLM latency and give
// the caller time to cancel the request mid-cascade.
// - Nothing here takes fc.mx or reaches into a worker, so the cascade cannot
// deadlock against, or stall, any other flow operation.
//
// A running flow picks the change up on the user's next input (which already
// re-resolves the provider by name and calls flowWorker.switchProvider) or on
// the next backend start (which rebuilds the provider from the DB row). Both
// paths compare the provider's raw configuration, so they also catch the case
// where the name did not change but the configuration behind it did.
//
// The two sweeps only match rows still bearing oldName, which makes the whole
// operation idempotent and safe to retry. They are issued independently and
// their errors are joined, so a failure on one table never silently skips the
// other.
func (fc *flowController) reassignFlowsProvider(
ctx context.Context,
userID int64,
oldName, newName provider.ProviderName,
) error {
logger := logrus.WithContext(ctx).WithFields(logrus.Fields{
"user_id": userID,
"old_name": oldName.String(),
"new_name": newName.String(),
})
if oldName == newName {
logger.Debug("provider name unchanged, nothing to reassign")
return nil
}
// Only references that would otherwise dangle get rewritten. oldName can
// still resolve after the provider is gone when it named an override of a
// built-in — an intentional feature — in which case the built-in answers to
// that name again and the stored value is already correct. Rewriting it
// anyway would repoint rows that predate the override, and (when the
// override's type differed from the built-in it was named after) would send
// them to the wrong default entirely.
if _, err := fc.provs.GetProvider(ctx, oldName, userID); err == nil {
logger.Debug("old provider name still resolves, nothing to reassign")
return nil
}
// Detached from the caller's request context: these are two short statements
// and the reference must not be left half-rewritten because a browser tab
// was closed. The timeout keeps a stuck DB from pinning the goroutine.
ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), reassignProviderTimeout)
defer cancel()
flows, flowsErr := fc.db.UpdateFlowsProviderNameByOldName(ctx, database.UpdateFlowsProviderNameByOldNameParams{
NewName: newName.String(),
UserID: userID,
OldName: oldName.String(),
})
if flowsErr != nil {
logger.WithError(flowsErr).Error("failed to bulk-update flows provider name")
flowsErr = fmt.Errorf("failed to bulk-update flows provider name: %w", flowsErr)
}
assistants, asstErr := fc.db.UpdateAssistantsProviderNameByOldName(
ctx, database.UpdateAssistantsProviderNameByOldNameParams{
NewName: newName.String(),
UserID: userID,
OldName: oldName.String(),
})
if asstErr != nil {
logger.WithError(asstErr).Error("failed to bulk-update assistants provider name")
asstErr = fmt.Errorf("failed to bulk-update assistants provider name: %w", asstErr)
}
// Publishing happens only after both writes are done. A subscriber that is
// not draining its channel makes each publish cost up to the subscription
// send timeout, so doing it in between would let a wedged websocket client
// eat the deadline and starve the second UPDATE.
for _, flow := range flows {
// Skipped rather than published with no containers: FlowUpdated carries
// the full terminal list and the client replaces its cached value with
// whatever arrives, so an empty list would wipe the flow's terminals in
// the UI. Same handling as flowWorker.switchProvider.
containers, err := fc.db.GetFlowContainers(ctx, flow.ID)
if err != nil {
logger.WithError(err).Warnf("failed to get containers for flow %d, skipping its update event", flow.ID)
continue
}
fc.subs.NewFlowPublisher(userID, flow.ID).FlowUpdated(ctx, flow, containers)
}
for _, assistant := range assistants {
fc.subs.NewFlowPublisher(userID, assistant.FlowID).AssistantUpdated(ctx, assistant)
}
logger.WithFields(logrus.Fields{
"flows_updated": len(flows),
"assistants_updated": len(assistants),
}).Info("provider reference reassigned")
return errors.Join(flowsErr, asstErr)
}
+288
View File
@@ -0,0 +1,288 @@
package controller
import (
"context"
"errors"
"fmt"
"sync"
"testing"
"pentagi/pkg/database"
"pentagi/pkg/graph/subscriptions"
"pentagi/pkg/providers"
"pentagi/pkg/providers/provider"
"pentagi/pkg/providers/tester/mock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// ---- fakes -----------------------------------------------------------------
// cascadeFakeQuerier embeds a nil database.Querier so any method the code under
// test calls without a stub here panics loudly (signalling a test gap) instead
// of silently returning zero values.
type cascadeFakeQuerier struct {
database.Querier
flowsCalls []database.UpdateFlowsProviderNameByOldNameParams
flowsResult []database.Flow
flowsErr error
assistantsCalls []database.UpdateAssistantsProviderNameByOldNameParams
assistantsResult []database.Assistant
assistantsErr error
containersResult []database.Container
containersErr error
}
func (f *cascadeFakeQuerier) UpdateFlowsProviderNameByOldName(
ctx context.Context, arg database.UpdateFlowsProviderNameByOldNameParams,
) ([]database.Flow, error) {
f.flowsCalls = append(f.flowsCalls, arg)
if f.flowsErr != nil {
return nil, f.flowsErr
}
return f.flowsResult, nil
}
func (f *cascadeFakeQuerier) UpdateAssistantsProviderNameByOldName(
ctx context.Context, arg database.UpdateAssistantsProviderNameByOldNameParams,
) ([]database.Assistant, error) {
f.assistantsCalls = append(f.assistantsCalls, arg)
if f.assistantsErr != nil {
return nil, f.assistantsErr
}
return f.assistantsResult, nil
}
func (f *cascadeFakeQuerier) GetFlowContainers(ctx context.Context, flowID int64) ([]database.Container, error) {
if f.containersErr != nil {
return nil, f.containersErr
}
return f.containersResult, nil
}
// cascadeFakePublisher records every FlowUpdated/AssistantUpdated call.
type cascadeFakePublisher struct {
subscriptions.FlowPublisher
mx sync.Mutex
flowUpdated []database.Flow
assistantUpdated []database.Assistant
}
func (p *cascadeFakePublisher) FlowUpdated(ctx context.Context, flow database.Flow, terms []database.Container) {
p.mx.Lock()
defer p.mx.Unlock()
p.flowUpdated = append(p.flowUpdated, flow)
}
func (p *cascadeFakePublisher) AssistantUpdated(ctx context.Context, assistant database.Assistant) {
p.mx.Lock()
defer p.mx.Unlock()
p.assistantUpdated = append(p.assistantUpdated, assistant)
}
// cascadeFakeSubscriptions always hands out the same publisher regardless of
// userID/flowID, so a test can assert on everything published in one place.
type cascadeFakeSubscriptions struct {
subscriptions.SubscriptionsController
pub *cascadeFakePublisher
}
func (s *cascadeFakeSubscriptions) NewFlowPublisher(userID, flowID int64) subscriptions.FlowPublisher {
return s.pub
}
// cascadeFakeProviders answers the "does the old name still resolve?" probe.
// resolvable lists the names that do; anything else fails, like a custom
// provider name that no built-in answers to.
type cascadeFakeProviders struct {
providers.ProviderController
resolvable map[provider.ProviderName]bool
}
func (p *cascadeFakeProviders) GetProvider(
ctx context.Context, prvname provider.ProviderName, userID int64,
) (provider.Provider, error) {
if p.resolvable[prvname] {
return mock.NewProvider(provider.ProviderQwen, prvname, "model"), nil
}
return nil, fmt.Errorf("provider not found by name '%s'", prvname)
}
// newTestFlowController builds a flowController with fake dependencies,
// bypassing NewFlowController (which needs a real docker client the provider
// reassignment does not touch). flows stays empty on purpose: the reassignment
// is DB-only and must never reach into a worker. resolvable names the providers
// that still answer after the change, which is what decides whether a rewrite
// is needed at all.
func newTestFlowController(
q *cascadeFakeQuerier, resolvable ...provider.ProviderName,
) (*flowController, *cascadeFakePublisher) {
pub := &cascadeFakePublisher{}
names := make(map[provider.ProviderName]bool, len(resolvable))
for _, name := range resolvable {
names[name] = true
}
return &flowController{
db: q,
mx: &sync.Mutex{},
flows: map[int64]FlowWorker{},
subs: &cascadeFakeSubscriptions{pub: pub},
provs: &cascadeFakeProviders{resolvable: names},
}, pub
}
// ---- tests ------------------------------------------------------------------
func TestReassignFlowsProvider_RenameSweepsFlowsAndAssistants(t *testing.T) {
const userID = int64(1)
q := &cascadeFakeQuerier{
flowsResult: []database.Flow{{ID: 10}, {ID: 11}},
assistantsResult: []database.Assistant{{ID: 100, FlowID: 10}},
}
fc, pub := newTestFlowController(q)
err := fc.RenameFlowsProvider(context.Background(), userID, "my-qwen", "my-qwen-renamed")
require.NoError(t, err)
require.Len(t, q.flowsCalls, 1)
assert.Equal(t, userID, q.flowsCalls[0].UserID, "the sweep must be scoped to the acting user")
assert.Equal(t, "my-qwen", q.flowsCalls[0].OldName)
assert.Equal(t, "my-qwen-renamed", q.flowsCalls[0].NewName)
require.Len(t, q.assistantsCalls, 1)
assert.Equal(t, userID, q.assistantsCalls[0].UserID)
assert.Equal(t, "my-qwen", q.assistantsCalls[0].OldName)
assert.Equal(t, "my-qwen-renamed", q.assistantsCalls[0].NewName)
assert.Len(t, pub.flowUpdated, 2, "every rewritten flow row must be published")
assert.Len(t, pub.assistantUpdated, 1, "every rewritten assistant row must be published")
}
func TestReassignFlowsProvider_ResetUsesProviderTypeAsDefaultName(t *testing.T) {
q := &cascadeFakeQuerier{}
fc, _ := newTestFlowController(q)
err := fc.ResetFlowsProviderToDefault(context.Background(), 1, "my-custom-qwen", provider.ProviderQwen)
require.NoError(t, err)
require.Len(t, q.flowsCalls, 1)
assert.Equal(t, "my-custom-qwen", q.flowsCalls[0].OldName)
assert.Equal(t, string(provider.ProviderQwen), q.flowsCalls[0].NewName,
"the default provider name is literally the provider type string")
}
// Deleting a user provider named exactly like the built-in of its type (an
// intentional way to override the product default) leaves the stored name
// already resolving to that built-in, so there is nothing to rewrite. Running
// flows drop the deleted configuration on their next input or on the next
// start — see flowProvider.SetProvider, which compares the raw config.
func TestReassignFlowsProvider_ShadowedDefaultNameIsNoOp(t *testing.T) {
q := &cascadeFakeQuerier{}
fc, pub := newTestFlowController(q, "qwen")
err := fc.ResetFlowsProviderToDefault(context.Background(), 1, "qwen", provider.ProviderQwen)
require.NoError(t, err)
assert.Empty(t, q.flowsCalls, "a no-op rename must not issue a self-matching UPDATE")
assert.Empty(t, q.assistantsCalls)
assert.Empty(t, pub.flowUpdated, "and must not republish every row of that provider")
}
// The override may also be named after a built-in of a DIFFERENT type (a
// "custom"-typed provider called "openai"). Deleting it must not drag those rows
// onto the "custom" built-in: the name "openai" resolves again on its own, and
// rewriting it would repoint rows that predate the override.
func TestReassignFlowsProvider_OldNameStillResolvesIsNoOp(t *testing.T) {
q := &cascadeFakeQuerier{}
fc, pub := newTestFlowController(q, "openai")
err := fc.ResetFlowsProviderToDefault(context.Background(), 1, "openai", provider.ProviderCustom)
require.NoError(t, err)
assert.Empty(t, q.flowsCalls, "a still-resolving name must never be rewritten")
assert.Empty(t, q.assistantsCalls)
assert.Empty(t, pub.flowUpdated)
}
// ...but if that built-in is not enabled, the name really would dangle, so the
// reset must go through and point the rows at the deleted provider's own type.
func TestReassignFlowsProvider_UnresolvableOldNameIsRewritten(t *testing.T) {
q := &cascadeFakeQuerier{}
fc, _ := newTestFlowController(q) // nothing resolves
err := fc.ResetFlowsProviderToDefault(context.Background(), 1, "openai", provider.ProviderCustom)
require.NoError(t, err)
require.Len(t, q.flowsCalls, 1)
assert.Equal(t, "openai", q.flowsCalls[0].OldName)
assert.Equal(t, string(provider.ProviderCustom), q.flowsCalls[0].NewName)
}
func TestReassignFlowsProvider_FlowsSweepErrorStillRunsAssistantsSweep(t *testing.T) {
flowsErr := errors.New("flows update exploded")
q := &cascadeFakeQuerier{
flowsErr: flowsErr,
assistantsResult: []database.Assistant{{ID: 100, FlowID: 10}},
}
fc, pub := newTestFlowController(q)
err := fc.RenameFlowsProvider(context.Background(), 1, "old", "new")
require.Error(t, err)
assert.ErrorIs(t, err, flowsErr, "the failure must be reported, not swallowed")
require.Len(t, q.assistantsCalls, 1,
"the tables are independent: a failure on one must not skip the other")
assert.Len(t, pub.assistantUpdated, 1)
}
func TestReassignFlowsProvider_BothSweepErrorsAreReported(t *testing.T) {
flowsErr, assistantsErr := errors.New("flows boom"), errors.New("assistants boom")
q := &cascadeFakeQuerier{flowsErr: flowsErr, assistantsErr: assistantsErr}
fc, _ := newTestFlowController(q)
err := fc.RenameFlowsProvider(context.Background(), 1, "old", "new")
require.Error(t, err)
assert.ErrorIs(t, err, flowsErr)
assert.ErrorIs(t, err, assistantsErr)
}
// FlowUpdated carries the flow's full terminal list and the client replaces its
// cached value wholesale, so publishing with no containers would blank the
// terminals in the UI. Better to skip the event than to corrupt the view.
func TestReassignFlowsProvider_SkipsPublishWhenContainerLookupFails(t *testing.T) {
q := &cascadeFakeQuerier{
flowsResult: []database.Flow{{ID: 10}},
containersErr: errors.New("containers unavailable"),
}
fc, pub := newTestFlowController(q)
err := fc.RenameFlowsProvider(context.Background(), 1, "old", "new")
require.NoError(t, err, "a container lookup failure is not a cascade failure")
assert.Empty(t, pub.flowUpdated, "publishing an empty terminal list would wipe the client's cache")
require.Len(t, q.flowsCalls, 1, "the rewrite itself must still have happened")
}
// The sweep is detached from the caller's context on purpose: the provider row
// is already committed by the time it runs, so a disconnecting HTTP client must
// not leave half of the references pointing at a name that no longer exists.
func TestReassignFlowsProvider_RunsDespiteCancelledCallerContext(t *testing.T) {
q := &cascadeFakeQuerier{flowsResult: []database.Flow{{ID: 10}}}
fc, _ := newTestFlowController(q)
ctx, cancel := context.WithCancel(context.Background())
cancel()
err := fc.RenameFlowsProvider(ctx, 1, "old", "new")
require.NoError(t, err)
assert.Len(t, q.flowsCalls, 1)
assert.Len(t, q.assistantsCalls, 1)
}
+63
View File
@@ -589,3 +589,66 @@ func (q *Queries) UpdateAssistantUseAgents(ctx context.Context, arg UpdateAssist
)
return i, err
}
const updateAssistantsProviderNameByOldName = `-- name: UpdateAssistantsProviderNameByOldName :many
UPDATE assistants a
SET model_provider_name = $1
FROM flows f
WHERE a.flow_id = f.id
AND f.user_id = $2
AND a.model_provider_name = $3
AND a.deleted_at IS NULL
AND f.deleted_at IS NULL
RETURNING a.id, a.status, a.title, a.model, a.model_provider_name, a.language, a.functions, a.trace_id, a.flow_id, a.use_agents, a.msgchain_id, a.created_at, a.updated_at, a.deleted_at, a.model_provider_type, a.tool_call_id_template
`
type UpdateAssistantsProviderNameByOldNameParams struct {
NewName string `json:"new_name"`
UserID int64 `json:"user_id"`
OldName string `json:"old_name"`
}
// The assistants counterpart of UpdateFlowsProviderNameByOldName. Assistants
// carry their own provider reference, independent of their flow's, and the
// table has no user_id — ownership is derived by joining through flows, so the
// statement can never cross a tenant boundary. Idempotent for the same reason:
// a rewritten row no longer matches old_name.
func (q *Queries) UpdateAssistantsProviderNameByOldName(ctx context.Context, arg UpdateAssistantsProviderNameByOldNameParams) ([]Assistant, error) {
rows, err := q.db.QueryContext(ctx, updateAssistantsProviderNameByOldName, arg.NewName, arg.UserID, arg.OldName)
if err != nil {
return nil, err
}
defer rows.Close()
var items []Assistant
for rows.Next() {
var i Assistant
if err := rows.Scan(
&i.ID,
&i.Status,
&i.Title,
&i.Model,
&i.ModelProviderName,
&i.Language,
&i.Functions,
&i.TraceID,
&i.FlowID,
&i.UseAgents,
&i.MsgchainID,
&i.CreatedAt,
&i.UpdatedAt,
&i.DeletedAt,
&i.ModelProviderType,
&i.ToolCallIDTemplate,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
+56
View File
@@ -700,3 +700,59 @@ func (q *Queries) UpdateFlowToolCallIDTemplate(ctx context.Context, arg UpdateFl
)
return i, err
}
const updateFlowsProviderNameByOldName = `-- name: UpdateFlowsProviderNameByOldName :many
UPDATE flows
SET model_provider_name = $1
WHERE user_id = $2 AND model_provider_name = $3 AND deleted_at IS NULL
RETURNING id, status, title, model, model_provider_name, language, functions, user_id, created_at, updated_at, deleted_at, trace_id, model_provider_type, tool_call_id_template
`
type UpdateFlowsProviderNameByOldNameParams struct {
NewName string `json:"new_name"`
UserID int64 `json:"user_id"`
OldName string `json:"old_name"`
}
// Bulk-renames every flow row of a user still pointing at a provider's old name
// (the user renamed a custom LLM provider, or deleted one and the reference is
// reset to the built-in name of its type). Matching on old_name makes the
// statement idempotent: once rewritten, a row no longer matches, so it is safe
// to call unconditionally and to retry.
func (q *Queries) UpdateFlowsProviderNameByOldName(ctx context.Context, arg UpdateFlowsProviderNameByOldNameParams) ([]Flow, error) {
rows, err := q.db.QueryContext(ctx, updateFlowsProviderNameByOldName, arg.NewName, arg.UserID, arg.OldName)
if err != nil {
return nil, err
}
defer rows.Close()
var items []Flow
for rows.Next() {
var i Flow
if err := rows.Scan(
&i.ID,
&i.Status,
&i.Title,
&i.Model,
&i.ModelProviderName,
&i.Language,
&i.Functions,
&i.UserID,
&i.CreatedAt,
&i.UpdatedAt,
&i.DeletedAt,
&i.TraceID,
&i.ModelProviderType,
&i.ToolCallIDTemplate,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
+7 -1
View File
@@ -76,7 +76,7 @@ func (q *Queries) DeleteProvider(ctx context.Context, id int64) (Provider, error
const deleteUserProvider = `-- name: DeleteUserProvider :one
UPDATE providers
SET deleted_at = CURRENT_TIMESTAMP
WHERE id = $1 AND user_id = $2
WHERE id = $1 AND user_id = $2 AND deleted_at IS NULL
RETURNING id, user_id, type, name, config, created_at, updated_at, deleted_at
`
@@ -85,6 +85,12 @@ type DeleteUserProviderParams struct {
UserID int64 `json:"user_id"`
}
// deleted_at IS NULL is load-bearing, not just hygiene: without it a replayed
// delete of an already-deleted id still succeeds and still returns the row, and
// the caller then resets every flow/assistant matching that name — which by
// then may belong to a *different*, live provider reusing the freed name
// (providers_name_user_id_unique is partial on deleted_at IS NULL). Mirrors the
// guard GetUserProvider already applies on the rename path.
func (q *Queries) DeleteUserProvider(ctx context.Context, arg DeleteUserProviderParams) (Provider, error) {
row := q.db.QueryRowContext(ctx, deleteUserProvider, arg.ID, arg.UserID)
var i Provider
+18
View File
@@ -57,6 +57,12 @@ type Querier interface {
DeleteUserKnowledgeDocument(ctx context.Context, arg DeleteUserKnowledgeDocumentParams) error
DeleteUserPreferences(ctx context.Context, userID int64) error
DeleteUserPrompt(ctx context.Context, arg DeleteUserPromptParams) error
// deleted_at IS NULL is load-bearing, not just hygiene: without it a replayed
// delete of an already-deleted id still succeeds and still returns the row, and
// the caller then resets every flow/assistant matching that name — which by
// then may belong to a *different*, live provider reusing the freed name
// (providers_name_user_id_unique is partial on deleted_at IS NULL). Mirrors the
// guard GetUserProvider already applies on the rename path.
DeleteUserProvider(ctx context.Context, arg DeleteUserProviderParams) (Provider, error)
GetAPIToken(ctx context.Context, id int64) (ApiToken, error)
GetAPITokenByTokenID(ctx context.Context, tokenID string) (ApiToken, error)
@@ -272,6 +278,12 @@ type Querier interface {
UpdateAssistantTitle(ctx context.Context, arg UpdateAssistantTitleParams) (Assistant, error)
UpdateAssistantToolCallIDTemplate(ctx context.Context, arg UpdateAssistantToolCallIDTemplateParams) (Assistant, error)
UpdateAssistantUseAgents(ctx context.Context, arg UpdateAssistantUseAgentsParams) (Assistant, error)
// The assistants counterpart of UpdateFlowsProviderNameByOldName. Assistants
// carry their own provider reference, independent of their flow's, and the
// table has no user_id — ownership is derived by joining through flows, so the
// statement can never cross a tenant boundary. Idempotent for the same reason:
// a rewritten row no longer matches old_name.
UpdateAssistantsProviderNameByOldName(ctx context.Context, arg UpdateAssistantsProviderNameByOldNameParams) ([]Assistant, error)
UpdateContainerImage(ctx context.Context, arg UpdateContainerImageParams) (Container, error)
UpdateContainerStatus(ctx context.Context, arg UpdateContainerStatusParams) (Container, error)
UpdateContainerStatusLocalID(ctx context.Context, arg UpdateContainerStatusLocalIDParams) (Container, error)
@@ -282,6 +294,12 @@ type Querier interface {
UpdateFlowTemplate(ctx context.Context, arg UpdateFlowTemplateParams) (FlowTemplate, error)
UpdateFlowTitle(ctx context.Context, arg UpdateFlowTitleParams) (Flow, error)
UpdateFlowToolCallIDTemplate(ctx context.Context, arg UpdateFlowToolCallIDTemplateParams) (Flow, error)
// Bulk-renames every flow row of a user still pointing at a provider's old name
// (the user renamed a custom LLM provider, or deleted one and the reference is
// reset to the built-in name of its type). Matching on old_name makes the
// statement idempotent: once rewritten, a row no longer matches, so it is safe
// to call unconditionally and to retry.
UpdateFlowsProviderNameByOldName(ctx context.Context, arg UpdateFlowsProviderNameByOldNameParams) ([]Flow, error)
// Update an existing document's embedding, text and metadata atomically.
// embedding must be formatted as a PostgreSQL vector literal: '[f1,f2,...]'
// cmetadata must be valid JSON text.
+35 -5
View File
@@ -201,11 +201,12 @@ func (r *mutationResolver) DeleteFlow(ctx context.Context, flowID int64) (model.
"flow": flowID,
}).Debug("delete flow")
if fw, err := r.Controller.GetFlow(ctx, flowID); err == nil {
if err := fw.Finish(ctx); err != nil {
return model.ResultTypeError, err
}
} else if !errors.Is(err, controller.ErrFlowNotFound) {
// Goes through the controller rather than GetFlow + worker.Finish: only
// FinishFlow evicts the worker from the in-memory map. A worker left behind
// for a soft-deleted flow leaks for the process lifetime, still shows up in
// ListFlows, and can be finished a second time by a concurrent caller.
if err := r.Controller.FinishFlow(ctx, flowID); err != nil &&
!errors.Is(err, controller.ErrFlowNotFound) {
return model.ResultTypeError, err
}
@@ -555,6 +556,15 @@ func (r *mutationResolver) UpdateProvider(ctx context.Context, providerID int64,
"name": name,
}).Debug("update provider")
// Fetch the current name before the rename so the flow/assistant cascade
// below knows which old name to look for; UpdateProvider itself only
// returns the row *after* the rename.
existing, err := r.DB.GetUserProvider(ctx, database.GetUserProviderParams{ID: providerID, UserID: uid})
if err != nil {
return nil, fmt.Errorf("failed to get provider %d: %w", providerID, err)
}
oldName := provider.ProviderName(existing.Name)
cfg := converter.ConvertAgentsConfigFromGqlModel(&agents)
prvname := provider.ProviderName(name)
prv, err := r.ProvidersCtrl.UpdateProvider(ctx, uid, providerID, prvname, cfg)
@@ -564,6 +574,15 @@ func (r *mutationResolver) UpdateProvider(ctx context.Context, providerID int64,
r.Subscriptions.NewProviderPublisher(uid).ProviderUpdated(ctx, prv, cfg)
// Repoint every flow/assistant that referred to the old name so they keep
// resolving to a valid provider. This must not fail the mutation: the
// provider itself is already renamed, and the sweep is idempotent.
if oldName != prvname {
if err := r.Controller.RenameFlowsProvider(ctx, uid, oldName, prvname); err != nil {
r.Logger.WithError(err).Error("failed to cascade provider rename to flows/assistants")
}
}
return converter.ConvertProvider(prv, cfg), nil
}
@@ -584,6 +603,17 @@ func (r *mutationResolver) DeleteProvider(ctx context.Context, providerID int64)
return model.ResultTypeError, err
}
// Runs before anything else that can fail: the provider row is already
// soft-deleted, so bailing out earlier would leave the references dangling.
// Rows whose stored name still resolves (an override of a built-in) are left
// alone — running flows drop the deleted configuration on their next input or
// on the next start. Never fails the mutation: the provider is gone either way.
deletedName := provider.ProviderName(prv.Name)
deletedType := provider.ProviderType(prv.Type)
if err := r.Controller.ResetFlowsProviderToDefault(ctx, uid, deletedName, deletedType); err != nil {
r.Logger.WithError(err).Error("failed to cascade provider deletion to flows/assistants")
}
var cfg pconfig.ProviderConfig
if err := json.Unmarshal(prv.Config, &cfg); err != nil {
return model.ResultTypeError, err
+42 -9
View File
@@ -1,6 +1,7 @@
package providers
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -90,7 +91,11 @@ type FlowProvider interface {
SetTitle(title string)
SetAgentLogProvider(agentLog tools.AgentLogProvider)
SetMsgLogProvider(msgLog tools.MsgLogProvider)
SetProvider(ctx context.Context, newProvider provider.Provider) error
// SetProvider swaps the provider instance backing this flow. It reports
// whether anything actually changed and, when it did, the tool call ID
// template resolved for the new provider — so the caller can persist a
// consistent (provider, template) pair without re-reading shared state.
SetProvider(ctx context.Context, newProvider provider.Provider) (bool, string, error)
GetTaskTitle(ctx context.Context, input string) (string, error)
GenerateSubtasks(ctx context.Context, taskID int64) ([]tools.SubtaskInfo, error)
@@ -180,22 +185,50 @@ func (fp *flowProvider) SetMsgLogProvider(msgLog tools.MsgLogProvider) {
fp.msgLog = msgLog
}
func (fp *flowProvider) SetProvider(ctx context.Context, newProvider provider.Provider) error {
// SetProvider installs newProvider as the one backing this flow and returns
// (changed, toolCallIDTemplate, error).
//
// It is a no-op — changed=false — when newProvider is effectively the one
// already in use. "Effectively" means same name *and* same raw configuration:
// comparing names alone is not enough, because a user provider may be named
// exactly like a built-in one (an intentional feature: it lets a user override
// the product's default configuration for a provider type from the UI). When
// such an override is deleted or renamed away, the name a flow refers to stays
// valid but now resolves to a different configuration, and that switch must go
// through.
func (fp *flowProvider) SetProvider(ctx context.Context, newProvider provider.Provider) (bool, string, error) {
ctx, span := obs.Observer.NewSpan(ctx, obs.SpanKindInternal, "providers.flowProvider.SetProvider")
defer span.End()
if newProvider == nil {
return false, "", fmt.Errorf("new provider is nil")
}
fp.mx.RLock()
current, prompter := fp.Provider, fp.prompter
fp.mx.RUnlock()
if current != nil && current.Name() == newProvider.Name() &&
bytes.Equal(current.GetRawConfig(), newProvider.GetRawConfig()) {
return false, fp.ToolCallIDTemplate(), nil
}
// Resolved outside the lock on purpose: on a cold cache this probes the
// provider with live LLM calls, and holding the write lock across it would
// stall every agent chain reading through this flowProvider. Doing it first
// also keeps the swap atomic — a failure here leaves the flow untouched.
tcIDTemplate, err := newProvider.GetToolCallIDTemplate(ctx, prompter)
if err != nil {
return false, "", fmt.Errorf("failed to get tool call ID template: %w", err)
}
fp.mx.Lock()
defer fp.mx.Unlock()
fp.Provider = newProvider
fp.tcIDTemplate = tcIDTemplate
var err error
fp.tcIDTemplate, err = newProvider.GetToolCallIDTemplate(ctx, fp.prompter)
if err != nil {
return fmt.Errorf("failed to get tool call ID template: %w", err)
}
return nil
return true, tcIDTemplate, nil
}
func (fp *flowProvider) ID() int64 {
+87
View File
@@ -0,0 +1,87 @@
package providers
import (
"context"
"sync"
"testing"
"pentagi/pkg/providers/provider"
"pentagi/pkg/providers/tester/mock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// newSwitchTestFlowProvider builds the minimal flowProvider SetProvider needs:
// the lock it guards its state with, the currently installed provider and the
// tool call ID template resolved for it. Everything else stays nil — the mock
// provider ignores the prompter when resolving a template.
func newSwitchTestFlowProvider(current provider.Provider, tcIDTemplate string) *flowProvider {
return &flowProvider{
mx: &sync.RWMutex{},
tcIDTemplate: tcIDTemplate,
Provider: current,
}
}
func TestFlowProviderSetProvider_SameNameAndConfigIsNoOp(t *testing.T) {
current := mock.NewProvider(provider.ProviderQwen, "qwen", "qwen-model")
incoming := mock.NewProvider(provider.ProviderQwen, "qwen", "qwen-model")
fp := newSwitchTestFlowProvider(current, "call_{r:24:b}")
changed, tcIDTemplate, err := fp.SetProvider(context.Background(), incoming)
require.NoError(t, err)
assert.False(t, changed, "an identical provider must not trigger a switch")
assert.Equal(t, "call_{r:24:b}", tcIDTemplate, "the resolved template must be reported unchanged")
assert.Same(t, current, fp.Provider, "the installed provider instance must be left alone")
}
// The regression this guards: a user provider may be named exactly like a
// built-in one (an intentional feature — it overrides the product default for
// that provider type). Deleting or renaming such an override leaves the name a
// flow refers to valid while the configuration behind it changes, so comparing
// names alone would silently keep the flow on the deleted configuration.
func TestFlowProviderSetProvider_SameNameDifferentConfigSwitches(t *testing.T) {
override := mock.NewProvider(provider.ProviderOpenAI, "openai", "gpt-x")
override.SetRawConfig([]byte(`{"base_url":"https://gateway.internal"}`))
builtin := mock.NewProvider(provider.ProviderOpenAI, "openai", "gpt-x")
builtin.SetRawConfig([]byte(`{"base_url":"https://api.openai.com"}`))
fp := newSwitchTestFlowProvider(override, "stale_template")
changed, tcIDTemplate, err := fp.SetProvider(context.Background(), builtin)
require.NoError(t, err)
assert.True(t, changed, "same name but different configuration must switch")
assert.Equal(t, "toolu_{r:24:b}", tcIDTemplate, "the template must be re-resolved for the new provider")
assert.Same(t, builtin, fp.Provider)
assert.Equal(t, "toolu_{r:24:b}", fp.ToolCallIDTemplate(), "the stored template must be refreshed too")
}
func TestFlowProviderSetProvider_DifferentNameSwitches(t *testing.T) {
current := mock.NewProvider(provider.ProviderQwen, "my-qwen", "qwen-model")
incoming := mock.NewProvider(provider.ProviderQwen, "qwen", "qwen-model")
fp := newSwitchTestFlowProvider(current, "stale_template")
changed, tcIDTemplate, err := fp.SetProvider(context.Background(), incoming)
require.NoError(t, err)
assert.True(t, changed)
assert.Equal(t, "toolu_{r:24:b}", tcIDTemplate)
assert.Same(t, incoming, fp.Provider)
}
func TestFlowProviderSetProvider_NilProviderIsRejected(t *testing.T) {
current := mock.NewProvider(provider.ProviderQwen, "qwen", "qwen-model")
fp := newSwitchTestFlowProvider(current, "call_{r:24:b}")
changed, _, err := fp.SetProvider(context.Background(), nil)
require.Error(t, err)
assert.False(t, changed)
assert.Same(t, current, fp.Provider, "a rejected switch must leave the flow untouched")
}
@@ -26,6 +26,7 @@ type Provider struct {
streamingDelay time.Duration
providerConfig *pconfig.ProviderConfig
models pconfig.ModelsConfig
rawConfig []byte
// sequence, when set via SetSequentialResponses, makes CallWithTools
// ignore content-based matching and return each response strictly in
@@ -338,9 +339,19 @@ func (p *Provider) CallWithExtraOptions(
// GetRawConfig implements provider.Provider
func (p *Provider) GetRawConfig() []byte {
if p.rawConfig != nil {
return p.rawConfig
}
return []byte(`{"mock": true}`)
}
// SetRawConfig overrides what GetRawConfig returns, so a test can build two
// providers that share a name but not a configuration — exactly what a user
// provider named like a built-in one produces once it is deleted or renamed.
func (p *Provider) SetRawConfig(raw []byte) {
p.rawConfig = raw
}
// GetProviderConfig implements provider.Provider
func (p *Provider) GetProviderConfig() *pconfig.ProviderConfig {
if p.providerConfig != nil {
+16
View File
@@ -91,6 +91,22 @@ SET tool_call_id_template = $1
WHERE id = $2
RETURNING *;
-- name: UpdateAssistantsProviderNameByOldName :many
-- The assistants counterpart of UpdateFlowsProviderNameByOldName. Assistants
-- carry their own provider reference, independent of their flow's, and the
-- table has no user_id — ownership is derived by joining through flows, so the
-- statement can never cross a tenant boundary. Idempotent for the same reason:
-- a rewritten row no longer matches old_name.
UPDATE assistants a
SET model_provider_name = sqlc.arg(new_name)
FROM flows f
WHERE a.flow_id = f.id
AND f.user_id = sqlc.arg(user_id)
AND a.model_provider_name = sqlc.arg(old_name)
AND a.deleted_at IS NULL
AND f.deleted_at IS NULL
RETURNING a.id, a.status, a.title, a.model, a.model_provider_name, a.language, a.functions, a.trace_id, a.flow_id, a.use_agents, a.msgchain_id, a.created_at, a.updated_at, a.deleted_at, a.model_provider_type, a.tool_call_id_template;
-- name: DeleteAssistant :one
UPDATE assistants
SET deleted_at = CURRENT_TIMESTAMP
+11
View File
@@ -71,6 +71,17 @@ SET model_provider_name = $1, model_provider_type = $2, tool_call_id_template =
WHERE id = $5
RETURNING *;
-- name: UpdateFlowsProviderNameByOldName :many
-- Bulk-renames every flow row of a user still pointing at a provider's old name
-- (the user renamed a custom LLM provider, or deleted one and the reference is
-- reset to the built-in name of its type). Matching on old_name makes the
-- statement idempotent: once rewritten, a row no longer matches, so it is safe
-- to call unconditionally and to retry.
UPDATE flows
SET model_provider_name = sqlc.arg(new_name)
WHERE user_id = sqlc.arg(user_id) AND model_provider_name = sqlc.arg(old_name) AND deleted_at IS NULL
RETURNING *;
-- name: DeleteFlow :one
UPDATE flows
SET deleted_at = CURRENT_TIMESTAMP
+7 -1
View File
@@ -78,7 +78,13 @@ WHERE id = $1
RETURNING *;
-- name: DeleteUserProvider :one
-- deleted_at IS NULL is load-bearing, not just hygiene: without it a replayed
-- delete of an already-deleted id still succeeds and still returns the row, and
-- the caller then resets every flow/assistant matching that name — which by
-- then may belong to a *different*, live provider reusing the freed name
-- (providers_name_user_id_unique is partial on deleted_at IS NULL). Mirrors the
-- guard GetUserProvider already applies on the rename path.
UPDATE providers
SET deleted_at = CURRENT_TIMESTAMP
WHERE id = $1 AND user_id = $2
WHERE id = $1 AND user_id = $2 AND deleted_at IS NULL
RETURNING *;
+26 -3
View File
@@ -1,9 +1,14 @@
import { useQuery } from '@apollo/client/react';
import { useQuery, useSubscription } from '@apollo/client/react';
import { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import type { Provider } from '@/models/provider';
import { ProvidersDocument } from '@/graphql/types';
import {
ProviderCreatedDocument,
ProviderDeletedDocument,
ProvidersDocument,
ProviderUpdatedDocument,
} from '@/graphql/types';
import { findProviderByName, sortProviders } from '@/models/provider';
import { useUser } from '@/providers/user-provider';
@@ -24,10 +29,28 @@ interface ProvidersProviderProps {
export function ProvidersProvider({ children }: ProvidersProviderProps) {
const { isAuthenticated } = useUser();
const { data: providersData } = useQuery(ProvidersDocument, {
const { data: providersData, refetch: refetchProviders } = useQuery(ProvidersDocument, {
skip: !isAuthenticated(),
});
// The providers list is a separate root field from the settings page's own
// query, so editing a provider there leaves this copy stale — and a stale
// copy means the composer offers a provider name the backend no longer
// resolves. Refetch on every provider mutation instead of waiting for a
// page reload.
const refetchOnProviderEvent = useCallback(() => {
if (!isAuthenticated()) {
return;
}
void refetchProviders();
}, [isAuthenticated, refetchProviders]);
const subscriptionSkip = !isAuthenticated();
useSubscription(ProviderCreatedDocument, { onData: refetchOnProviderEvent, skip: subscriptionSkip });
useSubscription(ProviderUpdatedDocument, { onData: refetchOnProviderEvent, skip: subscriptionSkip });
useSubscription(ProviderDeletedDocument, { onData: refetchOnProviderEvent, skip: subscriptionSkip });
const providers = useMemo(() => sortProviders(providersData?.providers || []), [providersData?.providers]);
const [selectedProviderName, setSelectedProviderName] = useState<null | string>(() => {