fix: Runner run with new runctx

Change-Id: If0ae01e3a7d396d2c3080b359eff656a3ed43d0f

Author: luohq-bytedance

adk/flow.go

@@ -319,6 +319,10 @@ func initRunCtx(ctx context.Context, agentName string, input *AgentInput) (conte
 	return context.WithValue(ctx, runCtxKey{}, runCtx), runCtx
 }
 
+func ctxWithNewRunCtx(ctx context.Context) context.Context {
+	return context.WithValue(ctx, runCtxKey{}, &runContext{session: newRunSession()})
+}
+
 func getSession(ctx context.Context) *runSession {
 	v := ctx.Value(runCtxKey{})
 
@@ -369,7 +373,7 @@ func genAgentInput(runCtx *runContext, agentName string) (*AgentInput, error) {
 	return input, nil
 }
 
-func (a *flowAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {
+func (a *flowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
 	agentName := a.Name(ctx)
 
 	ctx, runCtx := initRunCtx(ctx, agentName, input)
@@ -384,10 +388,10 @@ func (a *flowAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRunOpt
 	}
 
 	if wf, ok := a.Agent.(*workflowAgent); ok {
-		return wf.Run(ctx, input)
+		return wf.Run(ctx, input, opts...)
 	}
 
-	aIter := a.Agent.Run(ctx, input)
+	aIter := a.Agent.Run(ctx, input, opts...)
 
 	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
 
@@ -434,7 +438,7 @@ func (a *flowAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRunOpt
 				return
 			}
 
-			subAIter := agentToRun.Run(ctx, input)
+			subAIter := agentToRun.Run(ctx, input, opts...)
 			for {
 				subEvent, ok_ := subAIter.Next()
 				if !ok_ {

adk/runner.go

@@ -22,41 +22,35 @@ import (
 	"github.com/cloudwego/eino/schema"
 )
 
-type Runner struct{}
-
-type RunnerConfig struct{}
+type Runner struct {
+	enableStreaming bool
+}
 
-func NewRunner(_ context.Context, _ *RunnerConfig) *Runner {
-	return &Runner{}
+type RunnerConfig struct {
+	EnableStreaming bool
 }
 
-type RunOptions struct {
-	enableStreaming bool
+func NewRunner(_ context.Context, conf RunnerConfig) *Runner {
+	return &Runner{enableStreaming: conf.EnableStreaming}
 }
 
-type RunOption func(*RunOptions)
+func (r *Runner) Run(ctx context.Context, agent Agent, msgs []Message,
+	opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
 
-func WithEnableStreaming() RunOption {
-	return func(opts *RunOptions) {
-		opts.enableStreaming = true
-	}
-}
-func (r *Runner) Run(ctx context.Context, agent Agent, msgs []Message, opts ...RunOption) *AsyncIterator[*AgentEvent] {
 	fa := toFlowAgent(agent)
 
-	options := &RunOptions{}
-	for _, opt := range opts {
-		opt(options)
-	}
-
 	input := &AgentInput{
 		Msgs:            msgs,
-		EnableStreaming: options.enableStreaming,
+		EnableStreaming: r.enableStreaming,
 	}
 
-	return fa.Run(ctx, input)
+	ctx = ctxWithNewRunCtx(ctx)
+
+	return fa.Run(ctx, input, opts...)
 }
 
-func (r *Runner) Query(ctx context.Context, agent Agent, query string, opts ...RunOption) *AsyncIterator[*AgentEvent] {
+func (r *Runner) Query(ctx context.Context, agent Agent,
+	query string, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
+
 	return r.Run(ctx, agent, []Message{schema.UserMessage(query)}, opts...)
 }

adk/workflow.go

@@ -53,7 +53,7 @@ func (a *workflowAgent) Description(_ context.Context) string {
 	return a.description
 }
 
-func (a *workflowAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRunOption) *AsyncIterator[*AgentEvent] {
+func (a *workflowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRunOption) *AsyncIterator[*AgentEvent] {
 	iterator, generator := NewAsyncIteratorPair[*AgentEvent]()
 
 	go func() {
@@ -74,11 +74,11 @@ func (a *workflowAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRu
 		// Different workflow execution based on mode
 		switch a.mode {
 		case workflowAgentModeSequential:
-			a.runSequential(ctx, input, generator)
+			a.runSequential(ctx, input, generator, opts...)
 		case workflowAgentModeLoop:
-			a.runLoop(ctx, input, generator)
+			a.runLoop(ctx, input, generator, opts...)
 		case workflowAgentModeParallel:
-			a.runParallel(ctx, input, generator)
+			a.runParallel(ctx, input, generator, opts...)
 		default:
 			err = errors.New(fmt.Sprintf("unsupported workflow agent mode: %d", a.mode))
 		}
@@ -88,10 +88,10 @@ func (a *workflowAgent) Run(ctx context.Context, input *AgentInput, _ ...AgentRu
 }
 
 func (a *workflowAgent) runSequential(ctx context.Context, input *AgentInput,
-	generator *AsyncGenerator[*AgentEvent]) (exit bool) {
+	generator *AsyncGenerator[*AgentEvent], opts ...AgentRunOption) (exit bool) {
 
 	for _, subAgent := range a.subAgents {
-		subIterator := subAgent.Run(ctx, input)
+		subIterator := subAgent.Run(ctx, input, opts...)
 		for {
 			event, ok := subIterator.Next()
 			if !ok {
@@ -116,22 +116,26 @@ func (a *workflowAgent) runSequential(ctx context.Context, input *AgentInput,
 	return false
 }
 
-func (a *workflowAgent) runLoop(ctx context.Context, input *AgentInput, generator *AsyncGenerator[*AgentEvent]) {
+func (a *workflowAgent) runLoop(ctx context.Context, input *AgentInput,
+	generator *AsyncGenerator[*AgentEvent], opts ...AgentRunOption) {
+
 	if len(a.subAgents) == 0 {
 		return
 	}
 
 	var iterations int
 	for iterations < a.maxIterations || a.maxIterations == 0 {
 		iterations++
-		exit := a.runSequential(ctx, input, generator)
+		exit := a.runSequential(ctx, input, generator, opts...)
 		if exit {
 			return
 		}
 	}
 }
 
-func (a *workflowAgent) runParallel(ctx context.Context, input *AgentInput, generator *AsyncGenerator[*AgentEvent]) {
+func (a *workflowAgent) runParallel(ctx context.Context, input *AgentInput,
+	generator *AsyncGenerator[*AgentEvent], opts ...AgentRunOption) {
+
 	if len(a.subAgents) == 0 {
 		return
 	}
@@ -150,7 +154,7 @@ func (a *workflowAgent) runParallel(ctx context.Context, input *AgentInput, gene
 					wg.Done()
 				}()
 
-				iterator := agent.Run(ctx, input)
+				iterator := agent.Run(ctx, input, opts...)
 				for {
 					event, ok := iterator.Next()
 					if !ok {