mirror of
https://github.com/vxcontrol/pentagi.git
synced 2026-08-24 03:56:32 +00:00
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:
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 *;
|
||||
|
||||
@@ -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>(() => {
|
||||
|
||||
Reference in New Issue
Block a user