core/server/handlers/stream_loop.go

package handlers

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

	"github.com/gin-gonic/gin"

	"git.gobha.me/xcaliber/chat-switchboard/events"
	"git.gobha.me/xcaliber/chat-switchboard/providers"
	"git.gobha.me/xcaliber/chat-switchboard/tools"
)

// streamResult holds the accumulated output from a streaming completion
// with tool execution. Callers are responsible for persistence.
type streamResult struct {
	Content             string
	ToolActivity        []map[string]interface{}
	InputTokens         int
	OutputTokens        int
	CacheCreationTokens int
	CacheReadTokens     int
}

// streamWithToolLoop is the single, canonical streaming implementation.
// It handles SSE setup, multi-round tool execution, reasoning_content
// forwarding, and client notifications. Returns the accumulated content
// and tool activity for the caller to persist however it needs to.
//
// Used by:
//   - streamCompletion (normal chat — persists via persistMessage)
//   - Regenerate       (regen/edit — persists as sibling with tree metadata)
//
// If you add features here (new event types, new tool capabilities, new
// persistence fields), both callers get them automatically.
func streamWithToolLoop(
	c *gin.Context,
	provider providers.Provider,
	cfg providers.ProviderConfig,
	req *providers.CompletionRequest,
	model, userID, channelID, personaID, workspaceID string,
	hub *events.Hub,
) streamResult {
	// 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)
	flusher, _ := c.Writer.(http.Flusher)

	flush := func() {
		if flusher != nil {
			flusher.Flush()
		}
	}

	sendSSE := func(data string) {
		fmt.Fprintf(c.Writer, "data: %s\n\n", data)
		flush()
	}

	sendEvent := func(event, data string) {
		fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, data)
		flush()
	}

	var result streamResult

	for iteration := 0; iteration < maxToolIterations; iteration++ {
		ch, err := provider.StreamCompletion(c.Request.Context(), cfg, *req)
		if err != nil {
			sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(err.Error())))
			return result
		}

		var iterContent string
		var iterReasoning string
		var toolCalls []providers.ToolCall

		for event := range ch {
			if event.Error != nil {
				sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(event.Error.Error())))
				return result
			}

			// Stream reasoning deltas to client
			if event.Reasoning != "" {
				iterReasoning += event.Reasoning
				sendSSE(fmt.Sprintf(
					`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q}`,
					event.Reasoning, model))
			}

			// Stream text deltas to client
			if event.Delta != "" {
				iterContent += event.Delta
				sendSSE(fmt.Sprintf(
					`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q}`,
					event.Delta, model))
			}

			if event.Done {
				if event.FinishReason == "tool_calls" && len(event.ToolCalls) > 0 {
					toolCalls = event.ToolCalls
					// Accumulate tokens from this tool-call iteration
					result.InputTokens += event.InputTokens
					result.OutputTokens += event.OutputTokens
					result.CacheCreationTokens += event.CacheCreationTokens
					result.CacheReadTokens += event.CacheReadTokens
				} else {
					// Normal completion — send finish event
					finishReason := event.FinishReason
					if finishReason == "" {
						finishReason = "stop"
					}
					if iterReasoning != "" {
						result.Content += "<think>" + iterReasoning + "</think>"
					}
					result.Content += iterContent
					// Accumulate final tokens
					result.InputTokens += event.InputTokens
					result.OutputTokens += event.OutputTokens
					result.CacheCreationTokens += event.CacheCreationTokens
					result.CacheReadTokens += event.CacheReadTokens
					sendSSE(fmt.Sprintf(
						`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q}`,
						finishReason, model))
					sendSSE("[DONE]")
					return result
				}
				break
			}
		}

		// ── Tool execution round ────────────────
		if len(toolCalls) == 0 {
			// Stream ended without tool calls or finish — shouldn't happen
			sendSSE("[DONE]")
			if iterReasoning != "" {
				result.Content += "<think>" + iterReasoning + "</think>"
			}
			result.Content += iterContent
			return result
		}

		// Notify client about tool calls
		toolCallsJSON, _ := json.Marshal(toolCalls)
		sendEvent("tool_use", string(toolCallsJSON))

		// Append assistant message (with tool_calls, possibly with text) to conversation
		assistantMsg := providers.Message{
			Role:    "assistant",
			Content: iterContent,
		}
		for _, tc := range toolCalls {
			assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
		}
		req.Messages = append(req.Messages, assistantMsg)

		// Execute all tool calls
		execCtx := tools.ExecutionContext{
			UserID:      userID,
			ChannelID:   channelID,
			PersonaID:   personaID,
			WorkspaceID: workspaceID,
		}
		for _, tc := range toolCalls {
			call := tools.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			}

			var toolResult tools.ToolResult

			// Check if tool is registered locally (server-side)
			if tools.Get(call.Name) != nil {
				log.Printf("🔧 Executing tool (server): %s (call %s)", call.Name, call.ID)
				toolResult = tools.ExecuteCall(c.Request.Context(), execCtx, call)
			} else if hub != nil && hub.IsConnected(userID) {
				// Route to browser via WebSocket bridge
				log.Printf("🔧 Executing tool (browser): %s (call %s)", call.Name, call.ID)
				toolResult = executeBrowserTool(hub, userID, call)
			} else {
				log.Printf("⚠️ Unknown tool: %s (call %s)", call.Name, call.ID)
				toolResult = tools.ToolResult{
					ToolCallID: call.ID,
					Name:       call.Name,
					Content:    `{"error":"unknown tool or browser not connected"}`,
					IsError:    true,
				}
			}

			// Notify client about tool result
			resultJSON, _ := json.Marshal(map[string]interface{}{
				"tool_call_id": toolResult.ToolCallID,
				"name":         toolResult.Name,
				"content":      toolResult.Content,
				"is_error":     toolResult.IsError,
			})
			sendEvent("tool_result", string(resultJSON))

			// Collect for persistence
			result.ToolActivity = append(result.ToolActivity, map[string]interface{}{
				"id":        tc.ID,
				"name":      tc.Function.Name,
				"arguments": tc.Function.Arguments,
				"result":    toolResult.Content,
				"is_error":  toolResult.IsError,
			})

			// Append tool result to conversation for next iteration
			req.Messages = append(req.Messages, providers.Message{
				Role:       "tool",
				Content:    toolResult.Content,
				ToolCallID: toolResult.ToolCallID,
				Name:       toolResult.Name,
			})
		}

		result.Content += iterContent
		// Loop: send updated messages back to provider
	}

	// Hit max iterations
	log.Printf("⚠️ Tool loop hit max iterations (%d) for channel %s", maxToolIterations, channelID)
	sendSSE(fmt.Sprintf(
		`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q}`,
		escapeJSON("[Tool execution limit reached]"), model))
	sendSSE("[DONE]")

	return result
}

const browserToolTimeout = 30 * time.Second

// streamModelResponse is a variant of streamWithToolLoop used by multi-model
// fan-out (v0.20.0). It streams a single model's response within an already-
// established SSE connection — it does NOT set SSE headers or send [DONE].
//
// SSE deltas include a "model_display" field so the frontend can attribute
// each response to the correct model.
func streamModelResponse(
	c *gin.Context,
	provider providers.Provider,
	cfg providers.ProviderConfig,
	req *providers.CompletionRequest,
	model, displayName, userID, channelID, personaID, workspaceID string,
	hub *events.Hub,
) streamResult {
	flusher, _ := c.Writer.(http.Flusher)
	flush := func() {
		if flusher != nil {
			flusher.Flush()
		}
	}
	sendSSE := func(data string) {
		fmt.Fprintf(c.Writer, "data: %s\n\n", data)
		flush()
	}
	sendEvent := func(event, data string) {
		fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event, data)
		flush()
	}

	var result streamResult

	for iteration := 0; iteration < maxToolIterations; iteration++ {
		ch, err := provider.StreamCompletion(c.Request.Context(), cfg, *req)
		if err != nil {
			sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(err.Error())))
			return result
		}

		var iterContent string
		var iterReasoning string
		var toolCalls []providers.ToolCall

		for event := range ch {
			if event.Error != nil {
				sendSSE(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(event.Error.Error())))
				return result
			}

			if event.Reasoning != "" {
				iterReasoning += event.Reasoning
				sendSSE(fmt.Sprintf(
					`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q,"model_display":%q}`,
					event.Reasoning, model, displayName))
			}

			if event.Delta != "" {
				iterContent += event.Delta
				sendSSE(fmt.Sprintf(
					`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q,"model_display":%q}`,
					event.Delta, model, displayName))
			}

			if event.Done {
				if event.FinishReason == "tool_calls" && len(event.ToolCalls) > 0 {
					toolCalls = event.ToolCalls
					result.InputTokens += event.InputTokens
					result.OutputTokens += event.OutputTokens
					result.CacheCreationTokens += event.CacheCreationTokens
					result.CacheReadTokens += event.CacheReadTokens
				} else {
					finishReason := event.FinishReason
					if finishReason == "" {
						finishReason = "stop"
					}
					if iterReasoning != "" {
						result.Content += "<think>" + iterReasoning + "</think>"
					}
					result.Content += iterContent
					result.InputTokens += event.InputTokens
					result.OutputTokens += event.OutputTokens
					result.CacheCreationTokens += event.CacheCreationTokens
					result.CacheReadTokens += event.CacheReadTokens
					sendSSE(fmt.Sprintf(
						`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q,"model_display":%q}`,
						finishReason, model, displayName))
					// Do NOT send [DONE] — caller manages that for multi-model
					return result
				}
				break
			}
		}

		// Tool execution (same logic as streamWithToolLoop)
		if len(toolCalls) == 0 {
			if iterReasoning != "" {
				result.Content += "<think>" + iterReasoning + "</think>"
			}
			result.Content += iterContent
			return result
		}

		toolCallsJSON, _ := json.Marshal(toolCalls)
		sendEvent("tool_use", string(toolCallsJSON))

		assistantMsg := providers.Message{
			Role:    "assistant",
			Content: iterContent,
		}
		for _, tc := range toolCalls {
			assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
		}
		req.Messages = append(req.Messages, assistantMsg)

		execCtx := tools.ExecutionContext{
			UserID:      userID,
			ChannelID:   channelID,
			PersonaID:   personaID,
			WorkspaceID: workspaceID,
		}
		for _, tc := range toolCalls {
			call := tools.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			}

			var toolResult tools.ToolResult
			if tools.Get(call.Name) != nil {
				toolResult = tools.ExecuteCall(c.Request.Context(), execCtx, call)
			} else if hub != nil && hub.IsConnected(userID) {
				toolResult = executeBrowserTool(hub, userID, call)
			} else {
				toolResult = tools.ToolResult{
					ToolCallID: call.ID,
					Name:       call.Name,
					Content:    `{"error":"unknown tool or browser not connected"}`,
					IsError:    true,
				}
			}

			resultJSON, _ := json.Marshal(map[string]interface{}{
				"tool_call_id": toolResult.ToolCallID,
				"name":         toolResult.Name,
				"content":      toolResult.Content,
				"is_error":     toolResult.IsError,
			})
			sendEvent("tool_result", string(resultJSON))

			result.ToolActivity = append(result.ToolActivity, map[string]interface{}{
				"id":        tc.ID,
				"name":      tc.Function.Name,
				"arguments": tc.Function.Arguments,
				"result":    toolResult.Content,
				"is_error":  toolResult.IsError,
			})

			req.Messages = append(req.Messages, providers.Message{
				Role:       "tool",
				Content:    toolResult.Content,
				ToolCallID: toolResult.ToolCallID,
				Name:       toolResult.Name,
			})
		}

		result.Content += iterContent
	}

	// Hit max iterations
	log.Printf("⚠️ Multi-model tool loop hit max iterations (%d) for model %s", maxToolIterations, displayName)
	sendSSE(fmt.Sprintf(
		`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q,"model_display":%q}`,
		escapeJSON("[Tool execution limit reached]"), model, displayName))

	return result
}

// 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.SendToUser(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,
	}
}