@@ -26,6 +26,8 @@ import (
2626 "github.com/charmbracelet/crush/internal/shell"
2727)
2828
29+const streamChunkTimeout = 80 * time.Second
30+
2931type AgentEventType string
3032
3133const (
@@ -553,16 +555,25 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
553555 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
554556
555557 // Process each event in the stream.
556− for event := range eventChan {
557− if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
558− if errors.Is(processErr, context.Canceled) {
559− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
560− } else {
561− a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "API Error", processErr.Error())
558+loop:
559+ for {
560+ select {
561+ case event, ok := <-eventChan:
562+ if !ok {
563+ break loop
562564 }
563− return assistantMsg, nil, processErr
564− }
565− if ctx.Err() != nil {
565+ if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
566+ if errors.Is(processErr, context.Canceled) {
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():
566577 a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
567578 return assistantMsg, nil, ctx.Err()
568579 }