From db68cde21c82ee5ed22b7c8c04468c01a6dde684 Mon Sep 17 00:00:00 2001 From: Dmitry Ng <19asdek91@gmail.com> Date: Tue, 4 Aug 2026 01:06:58 +0300 Subject: [PATCH] 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. --- backend/pkg/controller/flow.go | 55 +++- backend/pkg/controller/flows.go | 141 +++++++++ backend/pkg/controller/flows_test.go | 288 ++++++++++++++++++ backend/pkg/database/assistants.sql.go | 63 ++++ backend/pkg/database/flows.sql.go | 56 ++++ backend/pkg/database/providers.sql.go | 8 +- backend/pkg/database/querier.go | 18 ++ backend/pkg/graph/schema.resolvers.go | 40 ++- backend/pkg/providers/provider.go | 51 +++- backend/pkg/providers/provider_test.go | 87 ++++++ backend/pkg/providers/tester/mock/provider.go | 11 + backend/sqlc/models/assistants.sql | 16 + backend/sqlc/models/flows.sql | 11 + backend/sqlc/models/providers.sql | 8 +- frontend/src/providers/providers-provider.tsx | 29 +- 15 files changed, 849 insertions(+), 33 deletions(-) create mode 100644 backend/pkg/controller/flows_test.go create mode 100644 backend/pkg/providers/provider_test.go diff --git a/backend/pkg/controller/flow.go b/backend/pkg/controller/flow.go index b5b83545..3ac25518 100644 --- a/backend/pkg/controller/flow.go +++ b/backend/pkg/controller/flow.go @@ -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 { diff --git a/backend/pkg/controller/flows.go b/backend/pkg/controller/flows.go index 6793e1c4..5326e3b7 100644 --- a/backend/pkg/controller/flows.go +++ b/backend/pkg/controller/flows.go @@ -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) +} diff --git a/backend/pkg/controller/flows_test.go b/backend/pkg/controller/flows_test.go new file mode 100644 index 00000000..16b38329 --- /dev/null +++ b/backend/pkg/controller/flows_test.go @@ -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) +} diff --git a/backend/pkg/database/assistants.sql.go b/backend/pkg/database/assistants.sql.go index 08cf4f03..05d04955 100644 --- a/backend/pkg/database/assistants.sql.go +++ b/backend/pkg/database/assistants.sql.go @@ -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 +} diff --git a/backend/pkg/database/flows.sql.go b/backend/pkg/database/flows.sql.go index e04a2132..aeb3cd16 100644 --- a/backend/pkg/database/flows.sql.go +++ b/backend/pkg/database/flows.sql.go @@ -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 +} diff --git a/backend/pkg/database/providers.sql.go b/backend/pkg/database/providers.sql.go index bd0384a8..027cf13e 100644 --- a/backend/pkg/database/providers.sql.go +++ b/backend/pkg/database/providers.sql.go @@ -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 diff --git a/backend/pkg/database/querier.go b/backend/pkg/database/querier.go index 2f065052..7cb656e1 100644 --- a/backend/pkg/database/querier.go +++ b/backend/pkg/database/querier.go @@ -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. diff --git a/backend/pkg/graph/schema.resolvers.go b/backend/pkg/graph/schema.resolvers.go index c37a64f9..90a84c95 100644 --- a/backend/pkg/graph/schema.resolvers.go +++ b/backend/pkg/graph/schema.resolvers.go @@ -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 diff --git a/backend/pkg/providers/provider.go b/backend/pkg/providers/provider.go index 019faada..0065004e 100644 --- a/backend/pkg/providers/provider.go +++ b/backend/pkg/providers/provider.go @@ -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 { diff --git a/backend/pkg/providers/provider_test.go b/backend/pkg/providers/provider_test.go new file mode 100644 index 00000000..7188bdb6 --- /dev/null +++ b/backend/pkg/providers/provider_test.go @@ -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") +} diff --git a/backend/pkg/providers/tester/mock/provider.go b/backend/pkg/providers/tester/mock/provider.go index e2e8ca2e..ceb47241 100644 --- a/backend/pkg/providers/tester/mock/provider.go +++ b/backend/pkg/providers/tester/mock/provider.go @@ -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 { diff --git a/backend/sqlc/models/assistants.sql b/backend/sqlc/models/assistants.sql index 652418cf..a93b646c 100644 --- a/backend/sqlc/models/assistants.sql +++ b/backend/sqlc/models/assistants.sql @@ -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 diff --git a/backend/sqlc/models/flows.sql b/backend/sqlc/models/flows.sql index 8767c8f4..95969270 100644 --- a/backend/sqlc/models/flows.sql +++ b/backend/sqlc/models/flows.sql @@ -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 diff --git a/backend/sqlc/models/providers.sql b/backend/sqlc/models/providers.sql index 394ab4ab..95818779 100644 --- a/backend/sqlc/models/providers.sql +++ b/backend/sqlc/models/providers.sql @@ -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 *; diff --git a/frontend/src/providers/providers-provider.tsx b/frontend/src/providers/providers-provider.tsx index df025012..d1c5dd4a 100644 --- a/frontend/src/providers/providers-provider.tsx +++ b/frontend/src/providers/providers-provider.tsx @@ -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(() => {