core/server/handlers/tool_loop.go

// Package handlers — tool_loop.go
//
// v0.27.2: Extracted core tool loop with sink abstraction.
//
// coreToolLoop is the single implementation of multi-round LLM completion
// with tool execution. All callers (SSE streaming, multi-model streaming,
// sync JSON, headless scheduler) provide a LoopSink for I/O and a LoopConfig
// for dependencies. The loop handles provider calls, tool dispatch (server +
// browser bridge), health recording, budget enforcement, and workspace events.
//
// Prior to this extraction, the same logic was duplicated in three places:
//   - streamWithToolLoop      (SSE streaming)
//   - streamModelResponse     (multi-model SSE)
//   - syncCompletion          (non-streaming JSON)
//
// Those functions are now thin wrappers over coreToolLoop.
package handlers

import (
	"context"
	"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"
)

// ── Types ──────────────────────────────────────

// defaultMaxRounds is the safety limit on tool call iterations.
// Callers can override via LoopBudget.MaxRounds.
const defaultMaxRounds = 10

// LoopResult holds the accumulated output from a completion with tool
// execution. Callers are responsible for persistence.
type LoopResult struct {
	Content             string
	ToolActivity        []map[string]interface{}
	InputTokens         int
	OutputTokens        int
	CacheCreationTokens int
	CacheReadTokens     int
	ToolCallCount       int    // total tool calls across all rounds
	BudgetExceeded      string // "" | "tokens" | "tool_calls" | "max_rounds"
	Error               error  // non-nil if the loop terminated due to a provider error
}

// LoopBudget controls resource limits for the tool loop.
// Zero values mean "use default" for MaxRounds and "unlimited" for the rest.
type LoopBudget struct {
	MaxRounds    int // max tool-call iterations; 0 = defaultMaxRounds
	MaxToolCalls int // total tool calls across all rounds; 0 = unlimited
	MaxTokens    int // total input+output tokens; 0 = unlimited
}

func (b LoopBudget) maxRounds() int {
	if b.MaxRounds > 0 {
		return b.MaxRounds
	}
	return defaultMaxRounds
}

// LoopConfig bundles all dependencies for the core tool loop.
type LoopConfig struct {
	Provider     providers.Provider
	Cfg          providers.ProviderConfig
	Req          *providers.CompletionRequest
	Model        string
	ProviderType string
	ExecCtx      tools.ExecutionContext
	Hub          *events.Hub    // nil for headless (scheduler)
	Health       HealthRecorder // nil for headless
	ConfigID     string
	Budget       LoopBudget
	Streaming    bool // true = StreamCompletion, false = ChatCompletion
}

// ── Sink Interface ─────────────────────────────

// LoopSink receives events from the core tool loop for I/O.
// Implementations control what happens with each event: write SSE to
// a gin.Context, accumulate silently, log to stdout, etc.
type LoopSink interface {
	// OnDelta receives a text content delta from the LLM.
	OnDelta(delta string)

	// OnReasoning receives a reasoning/thinking delta from the LLM.
	OnReasoning(delta string)

	// OnError receives a fatal error (provider failure, stream error).
	OnError(err error)

	// OnToolUse is called when the LLM requests tool calls.
	OnToolUse(calls []providers.ToolCall)

	// OnToolResult is called after each tool execution completes.
	OnToolResult(result map[string]interface{})

	// OnFinish is called when the LLM produces a finish_reason.
	OnFinish(reason string)

	// OnDone signals end of the entire completion (after finish or budget breach).
	OnDone()

	// OnMaxRounds is called when the loop hits its iteration limit.
	OnMaxRounds()
}

// ── Core Tool Loop ─────────────────────────────

// coreToolLoop is the single, canonical tool-loop implementation.
// It drives multi-round LLM completion with tool execution, budget
// enforcement, and health recording. I/O is delegated to the sink.
//
// The caller is responsible for:
//   - Setting SSE headers (if streaming to a client)
//   - Calling providers.GetHooks().PreRequest() before entry
//   - Persisting the result after return
//   - Sending [DONE] or HTTP response after return (via sink.OnDone)
func CoreToolLoop(ctx context.Context, lcfg LoopConfig, sink LoopSink) LoopResult {
	var result LoopResult
	maxRounds := lcfg.Budget.maxRounds()

	for iteration := 0; iteration < maxRounds; iteration++ {
		var iterContent string
		var iterReasoning string
		var toolCalls []providers.ToolCall
		var iterInput, iterOutput, iterCacheCreate, iterCacheRead int

		if lcfg.Streaming {
			done := runStreamingRound(ctx, lcfg, sink, &result,
				&iterContent, &iterReasoning, &toolCalls,
				&iterInput, &iterOutput, &iterCacheCreate, &iterCacheRead)
			if done {
				return result
			}
		} else {
			done := runSyncRound(ctx, lcfg, sink, &result,
				&iterContent, &iterReasoning, &toolCalls,
				&iterInput, &iterOutput, &iterCacheCreate, &iterCacheRead)
			if done {
				return result
			}
		}

		// ── Tool execution round ────────────────
		if len(toolCalls) == 0 {
			// Stream/call ended without tool calls or finish — edge case
			if iterReasoning != "" {
				result.Content += "<think>" + iterReasoning + "</think>"
			}
			result.Content += iterContent
			sink.OnDone()
			return result
		}

		// Budget check: tool calls (pre-execution)
		result.ToolCallCount += len(toolCalls)
		if lcfg.Budget.MaxToolCalls > 0 && result.ToolCallCount > lcfg.Budget.MaxToolCalls {
			result.BudgetExceeded = "tool_calls"
			result.Content += iterContent
			sink.OnFinish("budget_exceeded")
			sink.OnDone()
			return result
		}

		sink.OnToolUse(toolCalls)

		// 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)
		}
		lcfg.Req.Messages = append(lcfg.Req.Messages, assistantMsg)

		// Execute all tool calls
		for _, tc := range toolCalls {
			call := tools.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			}

			toolResult := executeToolCall(ctx, lcfg, call)

			// Notify sink
			resultMap := map[string]interface{}{
				"tool_call_id": toolResult.ToolCallID,
				"name":         toolResult.Name,
				"content":      toolResult.Content,
				"is_error":     toolResult.IsError,
			}
			sink.OnToolResult(resultMap)

			// Emit workspace.file.changed for live editor updates (v0.21.5)
			emitWorkspaceEvent(lcfg, call, toolResult)

			// 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
			lcfg.Req.Messages = append(lcfg.Req.Messages, providers.Message{
				Role:       "tool",
				Content:    toolResult.Content,
				ToolCallID: toolResult.ToolCallID,
				Name:       toolResult.Name,
			})
		}

		result.Content += iterContent

		// Budget check: tokens (post-round)
		if lcfg.Budget.MaxTokens > 0 && (result.InputTokens+result.OutputTokens) > lcfg.Budget.MaxTokens {
			result.BudgetExceeded = "tokens"
			sink.OnFinish("budget_exceeded")
			sink.OnDone()
			return result
		}

		// Loop: send updated messages back to provider
	}

	// Hit max iterations
	result.BudgetExceeded = "max_rounds"
	sink.OnMaxRounds()
	sink.OnDone()
	return result
}

// ── Streaming Round ────────────────────────────

// runStreamingRound executes one provider StreamCompletion call and
// processes the event stream. Returns true if the loop should return
// (normal finish or error), false if tool execution should follow.
func runStreamingRound(
	ctx context.Context,
	lcfg LoopConfig,
	sink LoopSink,
	result *LoopResult,
	iterContent, iterReasoning *string,
	toolCalls *[]providers.ToolCall,
	iterInput, iterOutput, iterCacheCreate, iterCacheRead *int,
) bool {
	callStart := time.Now()
	ch, err := lcfg.Provider.StreamCompletion(ctx, lcfg.Cfg, *lcfg.Req)
	if err != nil {
		recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, err)
		result.Error = err
		sink.OnError(err)
		return true
	}

	for event := range ch {
		// Apply provider-specific post-processing hooks (v0.22.1)
		if hooks := providers.GetHooks(lcfg.ProviderType); hooks != nil {
			hooks.PostStreamEvent(lcfg.Cfg, &event)
		}

		if event.Error != nil {
			recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, event.Error)
			result.Error = event.Error
			sink.OnError(event.Error)
			return true
		}

		// Reasoning deltas
		if event.Reasoning != "" {
			*iterReasoning += event.Reasoning
			sink.OnReasoning(event.Reasoning)
		}

		// Content deltas
		if event.Delta != "" {
			*iterContent += event.Delta
			sink.OnDelta(event.Delta)
		}

		if event.Done {
			recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, nil)

			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
				return false // continue to tool execution
			}

			// Normal completion
			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

			sink.OnFinish(finishReason)
			sink.OnDone()
			return true
		}
	}

	// Channel closed without Done event — shouldn't happen
	return false
}

// ── Sync Round ─────────────────────────────────

// runSyncRound executes one provider ChatCompletion call. Returns true
// if the loop should return, false if tool execution should follow.
func runSyncRound(
	ctx context.Context,
	lcfg LoopConfig,
	sink LoopSink,
	result *LoopResult,
	iterContent, iterReasoning *string,
	toolCalls *[]providers.ToolCall,
	iterInput, iterOutput, iterCacheCreate, iterCacheRead *int,
) bool {
	callStart := time.Now()
	resp, err := lcfg.Provider.ChatCompletion(ctx, lcfg.Cfg, *lcfg.Req)
	recordHealthFn(lcfg.Health, lcfg.ConfigID, callStart, err)
	if err != nil {
		result.Error = err
		sink.OnError(err)
		return true
	}

	*iterContent = resp.Content
	*iterInput = resp.InputTokens
	*iterOutput = resp.OutputTokens
	*iterCacheCreate = resp.CacheCreationTokens
	*iterCacheRead = resp.CacheReadTokens

	result.InputTokens += resp.InputTokens
	result.OutputTokens += resp.OutputTokens
	result.CacheCreationTokens += resp.CacheCreationTokens
	result.CacheReadTokens += resp.CacheReadTokens

	if resp.FinishReason == "tool_calls" && len(resp.ToolCalls) > 0 {
		*toolCalls = resp.ToolCalls
		return false // continue to tool execution
	}

	// Normal completion — deliver full content as single delta
	result.Content += resp.Content
	sink.OnDelta(resp.Content)
	finishReason := resp.FinishReason
	if finishReason == "" {
		finishReason = "stop"
	}
	sink.OnFinish(finishReason)
	sink.OnDone()
	return true
}

// ── Tool Execution ─────────────────────────────

// executeToolCall dispatches a single tool call: server-side first,
// browser bridge fallback, unknown tool error as last resort.
func executeToolCall(ctx context.Context, lcfg LoopConfig, call tools.ToolCall) tools.ToolResult {
	toolStart := time.Now()

	// Server-side tool
	if tools.Get(call.Name) != nil {
		log.Printf("🔧 Executing tool (server): %s (call %s)", call.Name, call.ID)
		result := tools.ExecuteCall(ctx, lcfg.ExecCtx, call)
		if lcfg.Health != nil {
			toolLatency := int(time.Since(toolStart).Milliseconds())
			if result.IsError {
				lcfg.Health.RecordToolError(call.Name, toolLatency, result.Content)
			} else {
				lcfg.Health.RecordToolSuccess(call.Name, toolLatency)
			}
		}
		return result
	}

	// Browser bridge (only available when hub is connected)
	if lcfg.Hub != nil && lcfg.Hub.IsConnected(lcfg.ExecCtx.UserID) {
		log.Printf("🔧 Executing tool (browser): %s (call %s)", call.Name, call.ID)
		return executeBrowserTool(lcfg.Hub, lcfg.ExecCtx.UserID, call)
	}

	// Unknown tool
	log.Printf("⚠️ Unknown tool: %s (call %s)", call.Name, call.ID)
	return tools.ToolResult{
		ToolCallID: call.ID,
		Name:       call.Name,
		Content:    `{"error":"unknown tool or browser not connected"}`,
		IsError:    true,
	}
}

// emitWorkspaceEvent sends workspace.file.changed when a write tool succeeds.
func emitWorkspaceEvent(lcfg LoopConfig, call tools.ToolCall, result tools.ToolResult) {
	if lcfg.Hub == nil || result.IsError || lcfg.ExecCtx.WorkspaceID == "" {
		return
	}
	if call.Name != "workspace_write" && call.Name != "workspace_patch" {
		return
	}
	var toolArgs struct {
		Path string `json:"path"`
	}
	if json.Unmarshal([]byte(call.Arguments), &toolArgs) == nil && toolArgs.Path != "" {
		lcfg.Hub.SendToUser(lcfg.ExecCtx.UserID, events.Event{
			Label: "workspace.file.changed",
			Payload: events.MustJSON(map[string]string{
				"workspace_id": lcfg.ExecCtx.WorkspaceID,
				"path":         toolArgs.Path,
			}),
		})
	}
}

// ── Sink Implementations ───────────────────────

// sseSink writes SSE events to a gin.Context writer for interactive streaming.
type sseSink struct {
	w     http.ResponseWriter
	flush func()
	model string
}

func newSSESink(c *gin.Context, model string) *sseSink {
	flusher, _ := c.Writer.(http.Flusher)
	return &sseSink{
		w:     c.Writer,
		model: model,
		flush: func() {
			if flusher != nil {
				flusher.Flush()
			}
		},
	}
}

func (s *sseSink) sendData(data string) {
	fmt.Fprintf(s.w, "data: %s\n\n", data)
	s.flush()
}

func (s *sseSink) sendEvent(event, data string) {
	fmt.Fprintf(s.w, "event: %s\ndata: %s\n\n", event, data)
	s.flush()
}

func (s *sseSink) OnDelta(delta string) {
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q}`,
		delta, s.model))
}

func (s *sseSink) OnReasoning(delta string) {
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q}`,
		delta, s.model))
}

func (s *sseSink) OnError(err error) {
	s.sendData(fmt.Sprintf(`{"error":"%s"}`, escapeJSON(err.Error())))
}

func (s *sseSink) OnToolUse(calls []providers.ToolCall) {
	toolCallsJSON, _ := json.Marshal(calls)
	s.sendEvent("tool_use", string(toolCallsJSON))
}

func (s *sseSink) OnToolResult(result map[string]interface{}) {
	resultJSON, _ := json.Marshal(result)
	s.sendEvent("tool_result", string(resultJSON))
}

func (s *sseSink) OnFinish(reason string) {
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q}`,
		reason, s.model))
}

func (s *sseSink) OnDone() {
	s.sendData("[DONE]")
}

func (s *sseSink) OnMaxRounds() {
	log.Printf("⚠️ Tool loop hit max iterations for model %s", s.model)
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{"content":"%s"},"finish_reason":"stop"}],"model":%q}`,
		escapeJSON("[Tool execution limit reached]"), s.model))
}

// sseModelSink extends sseSink with model_display attribution for
// multi-model streaming. Does NOT send [DONE] — the multiModelStream
// caller sends it after all models finish.
type sseModelSink struct {
	sseSink
	displayName string
}

func newSSEModelSink(c *gin.Context, model, displayName string) *sseModelSink {
	base := newSSESink(c, model)
	return &sseModelSink{
		sseSink:     *base,
		displayName: displayName,
	}
}

func (s *sseModelSink) OnDelta(delta string) {
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{"content":%q},"finish_reason":null}],"model":%q,"model_display":%q}`,
		delta, s.model, s.displayName))
}

func (s *sseModelSink) OnReasoning(delta string) {
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{"reasoning_content":%q},"finish_reason":null}],"model":%q,"model_display":%q}`,
		delta, s.model, s.displayName))
}

func (s *sseModelSink) OnFinish(reason string) {
	s.sendData(fmt.Sprintf(
		`{"choices":[{"delta":{},"finish_reason":%q}],"model":%q,"model_display":%q}`,
		reason, s.model, s.displayName))
}

// OnDone is a no-op for multi-model — caller sends [DONE] after all models.
func (s *sseModelSink) OnDone() {}

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

// accumSink accumulates silently — used by syncCompletion where the
// caller formats the HTTP response after the loop returns.
type accumSink struct{}

func (accumSink) OnDelta(string)                        {}
func (accumSink) OnReasoning(string)                    {}
func (accumSink) OnError(error)                         {}
func (accumSink) OnToolUse([]providers.ToolCall)         {}
func (accumSink) OnToolResult(map[string]interface{})    {}
func (accumSink) OnFinish(string)                       {}
func (accumSink) OnDone()                               {}
func (accumSink) OnMaxRounds()                          {}

// HeadlessSink is used by the task scheduler — accumulates and logs.
// No client connection, no SSE output.
type HeadlessSink struct {
	taskID string
}

func NewHeadlessSink(taskID string) *HeadlessSink {
	return &HeadlessSink{taskID: taskID}
}

func (s *HeadlessSink) OnDelta(string)    {}
func (s *HeadlessSink) OnReasoning(string) {}

func (s *HeadlessSink) OnError(err error) {
	log.Printf("[task:%s] Error: %v", s.taskID, err)
}

func (s *HeadlessSink) OnToolUse(calls []providers.ToolCall) {
	names := make([]string, len(calls))
	for i, tc := range calls {
		names[i] = tc.Function.Name
	}
	log.Printf("[task:%s] Tool calls: %v", s.taskID, names)
}

func (s *HeadlessSink) OnToolResult(result map[string]interface{}) {
	name, _ := result["name"].(string)
	isErr, _ := result["is_error"].(bool)
	if isErr {
		log.Printf("[task:%s] Tool %s failed: %v", s.taskID, name, result["content"])
	}
}

func (s *HeadlessSink) OnFinish(reason string) {
	log.Printf("[task:%s] Finished: %s", s.taskID, reason)
}

func (s *HeadlessSink) OnDone() {}

func (s *HeadlessSink) OnMaxRounds() {
	log.Printf("[task:%s] Hit max tool iterations", s.taskID)
}