fix: workflow onEndWithStream handles stream asynchronously (#518)

This commit is contained in:
shentongmartin 2025-08-04 11:58:34 +08:00 committed by GitHub
parent 08d1f6bf67
commit 3fe4031531
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 38 additions and 37 deletions

View File

@ -496,48 +496,49 @@ func (w *WorkflowHandler) OnEndWithStreamOutput(ctx context.Context, info *callb
return ctx return ctx
} }
// consumes the stream synchronously because the Exit node has already processed this stream synchronously. safego.Go(ctx, func() {
defer output.Close() defer output.Close()
fullOutput := make(map[string]any) fullOutput := make(map[string]any)
for { for {
chunk, e := output.Recv() chunk, e := output.Recv()
if e != nil { if e != nil {
if e == io.EOF { if e == io.EOF {
break break
}
if _, ok := schema.GetSourceName(e); ok {
continue
}
logs.Errorf("workflow OnEndWithStreamOutput failed to receive stream output: %v", e)
_ = w.OnError(ctx, info, e)
return
} }
fullOutput, e = nodes.ConcatTwoMaps(fullOutput, chunk.(map[string]any))
if _, ok := schema.GetSourceName(e); ok { if e != nil {
continue logs.Errorf("failed to concat two maps: %v", e)
return
} }
logs.Errorf("workflow OnEndWithStreamOutput failed to receive stream output: %v", e)
_ = w.OnError(ctx, info, e)
return ctx
} }
fullOutput, e = nodes.ConcatTwoMaps(fullOutput, chunk.(map[string]any))
if e != nil {
logs.Errorf("failed to concat two maps: %v", e)
return ctx
}
}
c := GetExeCtx(ctx) c := GetExeCtx(ctx)
e := &Event{ e := &Event{
Type: WorkflowSuccess, Type: WorkflowSuccess,
Context: c, Context: c,
Duration: time.Since(time.UnixMilli(c.StartTime)), Duration: time.Since(time.UnixMilli(c.StartTime)),
Output: fullOutput, Output: fullOutput,
}
if c.TokenCollector != nil {
usage := c.TokenCollector.wait()
e.Token = &TokenInfo{
InputToken: int64(usage.PromptTokens),
OutputToken: int64(usage.CompletionTokens),
TotalToken: int64(usage.TotalTokens),
} }
}
w.ch <- e if c.TokenCollector != nil {
usage := c.TokenCollector.wait()
e.Token = &TokenInfo{
InputToken: int64(usage.PromptTokens),
OutputToken: int64(usage.CompletionTokens),
TotalToken: int64(usage.TotalTokens),
}
}
w.ch <- e
})
return ctx return ctx
} }