From 5dafd81a3f1fdfd6f48cda4c36c9fb9ca00f41f7 Mon Sep 17 00:00:00 2001 From: lvxinyu-1117 Date: Tue, 5 Aug 2025 11:26:51 +0800 Subject: [PATCH] fix: correct token calculation error caused by ctxcache init (#559) --- .../api/handler/coze/workflow_service_test.go | 6 +++++ backend/application/workflow/init.go | 3 +++ .../internal/compose/designate_option.go | 9 ------- .../domain/workflow/service/global_handler.go | 26 +++++++++++++++++++ 4 files changed, 35 insertions(+), 9 deletions(-) create mode 100644 backend/domain/workflow/service/global_handler.go diff --git a/backend/api/handler/coze/workflow_service_test.go b/backend/api/handler/coze/workflow_service_test.go index 72977e15..34b3301d 100644 --- a/backend/api/handler/coze/workflow_service_test.go +++ b/backend/api/handler/coze/workflow_service_test.go @@ -32,6 +32,7 @@ import ( "github.com/alicebob/miniredis/v2" "github.com/bytedance/mockey" + "github.com/cloudwego/eino/callbacks" model2 "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" "github.com/cloudwego/hertz/pkg/app" @@ -102,6 +103,11 @@ import ( "github.com/coze-dev/coze-studio/backend/types/errno" ) +func TestMain(m *testing.M) { + callbacks.AppendGlobalHandlers(service.GetTokenCallbackHandler()) + os.Exit(m.Run()) +} + type wfTestRunner struct { t *testing.T h *server.Hertz diff --git a/backend/application/workflow/init.go b/backend/application/workflow/init.go index 7e6d5a90..1e8e6de8 100644 --- a/backend/application/workflow/init.go +++ b/backend/application/workflow/init.go @@ -22,6 +22,7 @@ import ( "github.com/cloudwego/eino/compose" "gorm.io/gorm" + "github.com/cloudwego/eino/callbacks" "github.com/coze-dev/coze-studio/backend/application/internal" wfdatabase "github.com/coze-dev/coze-studio/backend/crossdomain/workflow/database" wfknowledge "github.com/coze-dev/coze-studio/backend/crossdomain/workflow/knowledge" @@ -43,6 +44,7 @@ import ( crosssearch "github.com/coze-dev/coze-studio/backend/domain/workflow/crossdomain/search" crossvariable "github.com/coze-dev/coze-studio/backend/domain/workflow/crossdomain/variable" "github.com/coze-dev/coze-studio/backend/domain/workflow/service" + workflowservice "github.com/coze-dev/coze-studio/backend/domain/workflow/service" "github.com/coze-dev/coze-studio/backend/infra/contract/cache" "github.com/coze-dev/coze-studio/backend/infra/contract/coderunner" "github.com/coze-dev/coze-studio/backend/infra/contract/idgen" @@ -89,6 +91,7 @@ func InitService(ctx context.Context, components *ServiceComponents) (*Applicati crossmodel.SetManager(wfmodel.NewModelManager(components.ModelManager, nil)) crosscode.SetCodeRunner(components.CodeRunner) crosssearch.SetNotifier(wfsearch.NewNotify(components.DomainNotifier)) + callbacks.AppendGlobalHandlers(workflowservice.GetTokenCallbackHandler()) SVC.DomainSVC = workflowDomainSVC SVC.ImageX = components.ImageX diff --git a/backend/domain/workflow/internal/compose/designate_option.go b/backend/domain/workflow/internal/compose/designate_option.go index 4f64fe09..fb039068 100644 --- a/backend/domain/workflow/internal/compose/designate_option.go +++ b/backend/domain/workflow/internal/compose/designate_option.go @@ -32,7 +32,6 @@ import ( "github.com/coze-dev/coze-studio/backend/domain/workflow/internal/execute" "github.com/coze-dev/coze-studio/backend/domain/workflow/internal/nodes" "github.com/coze-dev/coze-studio/backend/domain/workflow/internal/nodes/llm" - "github.com/coze-dev/coze-studio/backend/pkg/ctxcache" "github.com/coze-dev/coze-studio/backend/pkg/lang/ptr" ) @@ -47,8 +46,6 @@ func (r *WorkflowRunner) designateOptions(ctx context.Context) (context.Context, sw = r.streamWriter ) - const tokenCallbackKey = "token_callback_key" - if wb.AppID != nil && exeCfg.AppID == nil { exeCfg.AppID = wb.AppID } @@ -124,12 +121,6 @@ func (r *WorkflowRunner) designateOptions(ctx context.Context) (context.Context, opts = append(opts, einoCompose.WithCheckPointID(strconv.FormatInt(executeID, 10))) } - if !ctxcache.HasKey(ctx, tokenCallbackKey) { - opts = append(opts, einoCompose.WithCallbacks(execute.GetTokenCallbackHandler())) - ctx = ctxcache.Init(ctx) - ctxcache.Store(ctx, tokenCallbackKey, true) - } - return ctx, opts, nil } diff --git a/backend/domain/workflow/service/global_handler.go b/backend/domain/workflow/service/global_handler.go new file mode 100644 index 00000000..8367203f --- /dev/null +++ b/backend/domain/workflow/service/global_handler.go @@ -0,0 +1,26 @@ +/* + * Copyright 2025 coze-dev Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package service + +import ( + "github.com/cloudwego/eino/callbacks" + "github.com/coze-dev/coze-studio/backend/domain/workflow/internal/execute" +) + +func GetTokenCallbackHandler() callbacks.Handler { + return execute.GetTokenCallbackHandler() +}