core/server/handlers/stream_loop.go

package handlers

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"github.com/gin-gonic/gin"

	"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 string,
) 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,
		}
		for _, tc := range toolCalls {
			call := tools.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			}

			log.Printf("🔧 Executing tool: %s (call %s)", call.Name, call.ID)
			toolResult := tools.ExecuteCall(c.Request.Context(), execCtx, call)

			// 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
}