git.iko.soy/crush
on main help sign in
fix(stream): stream hang, add stream timeout (#1070)
a48d4e3Vadim Inshakov
internal/llm/agent/agent.go+20−9
compare with main
unifiedsplit
@@ -26,6 +26,8 @@ import (
26 "github.com/charmbracelet/crush/internal/shell"
26 "github.com/charmbracelet/crush/internal/shell"
27)
27)
28
28
29+const streamChunkTimeout = 80 * time.Second
30+
29type AgentEventType string
31type AgentEventType string
30
32
31const (
33const (
@@ -553,16 +555,25 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
553 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
555 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
554
556
555 // Process each event in the stream.
557 // Process each event in the stream.
558+loop:
559+ for {
560+ select {
556− for event := range eventChan {
561+ case event, ok := <-eventChan:
557− if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
562+ if !ok {
558− if errors.Is(processErr, context.Canceled) {
563+ break loop
559− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
560− } else {
561− a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "API Error", processErr.Error())
562 }
564 }
563− return assistantMsg, nil, processErr
565+ if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
564− }
566+ if errors.Is(processErr, context.Canceled) {
565− if ctx.Err() != nil {
567+ a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
568+ } else {
569+ a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "API Error", processErr.Error())
570+ }
571+ return assistantMsg, nil, processErr
572+ }
573+ case <-time.After(streamChunkTimeout):
574+ a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "Stream timeout", "No chunk received within timeout")
575+ return assistantMsg, nil, fmt.Errorf("stream chunk timeout")
576+ case <-ctx.Done():
566 a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
577 a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
567 return assistantMsg, nil, ctx.Err()
578 return assistantMsg, nil, ctx.Err()
568 }
579 }