fix: correct token calculation error caused by ctxcache init (#559)

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

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

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
 }
 

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()
+}