core/server/handlers/stream_loop.go

package handlers

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"chat-switchboard/events"
	"chat-switchboard/providers"
	"chat-switchboard/sandbox"
	"chat-switchboard/store"
	"chat-switchboard/tools"
)

// recordHealthFn records provider health from standalone streaming functions.
func recordHealthFn(hr HealthRecorder, configID string, start time.Time, err error) {
	if hr == nil || configID == "" {
		return
	}
	latencyMs := int(time.Since(start).Milliseconds())
	if err != nil {
		errMsg := err.Error()
		if strings.Contains(errMsg, "HTTP 429") || strings.Contains(errMsg, "rate limit") || strings.Contains(errMsg, "Too Many Requests") {
			hr.RecordRateLimit(configID, latencyMs, errMsg)
		} else {
			hr.RecordError(configID, latencyMs, errMsg)
		}
	} else {
		hr.RecordSuccess(configID, latencyMs)
	}
}

// streamWithToolLoop is the SSE streaming entry point for single-model
// completion with tool execution. Sets SSE headers, delegates to
// CoreToolLoop with an sseSink, and returns the accumulated result.
//
// Used by:
//   - streamCompletion (normal chat — persists via persistMessage)
//   - Regenerate       (regen/edit — persists as sibling with tree metadata)
func streamWithToolLoop(
	c *gin.Context,
	provider providers.Provider,
	cfg providers.ProviderConfig,
	req *providers.CompletionRequest,
	model, providerType, userID, channelID, personaID, workspaceID, configID, teamID string,
	hub *events.Hub,
	health HealthRecorder,
	runner *sandbox.Runner,
	extTools map[string]*store.PackageRegistration,
) LoopResult {
	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")
	c.Status(http.StatusOK)

	sink := newSSESink(c, model)

	return CoreToolLoop(c.Request.Context(), LoopConfig{
		Provider:     provider,
		Cfg:          cfg,
		Req:          req,
		Model:        model,
		ProviderType: providerType,
		ExecCtx: tools.ExecutionContext{
			UserID:      userID,
			ChannelID:   channelID,
			PersonaID:   personaID,
			WorkspaceID: workspaceID,
			TeamID:      teamID,
		},
		Hub:       hub,
		Health:    health,
		ConfigID:  configID,
		Budget:    LoopBudget{},
		Streaming: true,
		Runner:    runner,
		ExtTools:  extTools,
	}, sink)
}

const browserToolTimeout = 30 * time.Second

// streamModelResponse is the multi-model streaming variant. Runs on an
// already-established SSE connection — does NOT set SSE headers or send
// [DONE]. SSE deltas include a "model_display" field for frontend attribution.
func streamModelResponse(
	c *gin.Context,
	provider providers.Provider,
	cfg providers.ProviderConfig,
	req *providers.CompletionRequest,
	model, providerType, displayName, userID, channelID, personaID, workspaceID, configID, teamID string,
	hub *events.Hub,
	health HealthRecorder,
	runner *sandbox.Runner,
	extTools map[string]*store.PackageRegistration,
) LoopResult {
	sink := newSSEModelSink(c, model, displayName)

	return CoreToolLoop(c.Request.Context(), LoopConfig{
		Provider:     provider,
		Cfg:          cfg,
		Req:          req,
		Model:        model,
		ProviderType: providerType,
		ExecCtx: tools.ExecutionContext{
			UserID:      userID,
			ChannelID:   channelID,
			PersonaID:   personaID,
			WorkspaceID: workspaceID,
			TeamID:      teamID,
		},
		Hub:       hub,
		Health:    health,
		ConfigID:  configID,
		Budget:    LoopBudget{},
		Streaming: true,
		Runner:    runner,
		ExtTools:  extTools,
	}, sink)
}

// executeBrowserTool sends a tool call to the user's browser via WebSocket
// and blocks until a result is returned or the timeout elapses.
func executeBrowserTool(hub *events.Hub, userID string, call tools.ToolCall) tools.ToolResult {
	b := make([]byte, 16)
	rand.Read(b)
	callID := hex.EncodeToString(b)

	// Publish tool.call event to the user's browser
	payload, _ := json.Marshal(map[string]interface{}{
		"call_id":   callID,
		"tool":      call.Name,
		"arguments": json.RawMessage(call.Arguments),
	})

	hub.PublishToUser(userID, events.Event{
		Label:   "tool.call." + callID,
		Payload: payload,
		Ts:      time.Now().UnixMilli(),
	})

	// Block until browser sends tool.result.{callID} or timeout
	event, ok := hub.GetBus().WaitFor("tool.result."+callID, browserToolTimeout)
	if !ok {
		log.Printf("⚠️ Browser tool %s timed out after %v (call %s)", call.Name, browserToolTimeout, callID)
		return tools.ToolResult{
			ToolCallID: call.ID,
			Name:       call.Name,
			Content:    `{"error":"browser tool timed out (30s) — browser may be disconnected"}`,
			IsError:    true,
		}
	}

	// Parse result from the browser event payload
	var result struct {
		CallID string `json:"call_id"`
		Result string `json:"result"`
		Error  string `json:"error"`
	}
	if err := json.Unmarshal(event.Payload, &result); err != nil {
		return tools.ToolResult{
			ToolCallID: call.ID,
			Name:       call.Name,
			Content:    `{"error":"invalid result from browser tool"}`,
			IsError:    true,
		}
	}

	if result.Error != "" {
		return tools.ToolResult{
			ToolCallID: call.ID,
			Name:       call.Name,
			Content:    result.Error,
			IsError:    true,
		}
	}

	return tools.ToolResult{
		ToolCallID: call.ID,
		Name:       call.Name,
		Content:    result.Result,
	}
}