MuyueWorkspace/internal/api/handlers_shell_chat.go

package api

import (
	"context"
	"encoding/json"
	"net/http"
	"strings"

	"github.com/muyue/muyue/internal/agent"
	"github.com/muyue/muyue/internal/orchestrator"
)

const maxShellToolIterations = 10

type ShellChatRequest struct {
	Message   string   `json:"message"`
	Context   string   `json:"context,omitempty"`
	History   []string `json:"history,omitempty"`
	Cwd       string   `json:"cwd,omitempty"`
	Platform  string   `json:"platform,omitempty"`
	Stream    bool     `json:"stream"`
}

type ShellChatResponse struct {
	Content   string        `json:"content,omitempty"`
	ToolCalls []ToolCallInfo `json:"tool_calls,omitempty"`
	Error     string        `json:"error,omitempty"`
}

type ToolCallInfo struct {
	ID     string                 `json:"id"`
	Name   string                 `json:"name"`
	Args   map[string]interface{} `json:"args"`
	Result *toolResponseData      `json:"result,omitempty"`
	Error  string                 `json:"error,omitempty"`
}

func (s *Server) handleShellChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		writeError(w, "POST only", http.StatusMethodNotAllowed)
		return
	}

	var req ShellChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}

	if req.Message == "" {
		writeError(w, "message is required", http.StatusBadRequest)
		return
	}

	orb, err := orchestrator.New(s.config)
	if err != nil {
		writeError(w, err.Error(), http.StatusServiceUnavailable)
		return
	}

	orb.SetSystemPrompt(s.buildShellSystemPrompt(req))
	orb.SetTools(s.agentToolsJSON)

	if req.Stream {
		s.handleShellChatStream(w, orb, req)
	} else {
		s.handleShellChatNonStream(w, orb, req)
	}
}

func (s *Server) buildShellSystemPrompt(req ShellChatRequest) string {
	var sb strings.Builder

	sb.WriteString(`Tu es l'assistant Shell de Muyue. Tu as accès à un terminal et peux aider l'utilisateur avec:
- Exécuter des commandes shell
- Expliquer des erreurs de commandes
- Suggérer des commandes appropriées pour la tâche demandée
- Lire et explorer des fichiers
- Configurer l'environnement de développement

Tu peux appeler des outils pour exécuter des commandes, lire des fichiers, etc. Sois précis et concis dans tes réponses.

`)

	if req.Cwd != "" {
		sb.WriteString("Répertoire courant: " + req.Cwd + "\n")
	}
	if req.Platform != "" {
		sb.WriteString("Plateforme: " + req.Platform + "\n")
	}
	if req.Context != "" {
		sb.WriteString("\nContexte du terminal:\n" + req.Context + "\n")
	}
	if len(req.History) > 0 {
		sb.WriteString("\nDernières commandes exécutées:\n")
		for _, h := range req.History {
			sb.WriteString("  " + h + "\n")
		}
	}

	return sb.String()
}

func (s *Server) handleShellChatStream(w http.ResponseWriter, orb *orchestrator.Orchestrator, req ShellChatRequest) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(http.StatusOK)
	flusher, canFlush := w.(http.Flusher)

	writeSSE := func(data map[string]interface{}) {
		b, _ := json.Marshal(data)
		w.Write([]byte("data: " + string(b) + "\n\n"))
		if canFlush {
			flusher.Flush()
		}
	}

	ctx := context.Background()
	messages := []orchestrator.Message{
		{Role: "user", Content: req.Message},
	}

	var finalContent string
	var toolCalls []ToolCallInfo

	for i := 0; i < maxShellToolIterations; i++ {
		resp, err := orb.SendWithTools(messages)
		if err != nil {
			writeSSE(map[string]interface{}{"error": err.Error()})
			return
		}

		choice := resp.Choices[0]
		content := cleanThinkingTags(choice.Message.Content)

		if content != "" {
			for _, ch := range strings.Split(content, "") {
				writeSSE(map[string]interface{}{"content": ch})
			}
			finalContent = content
		}

		if len(choice.Message.ToolCalls) == 0 {
			break
		}

		assistantMsg := orchestrator.Message{
			Role:      "assistant",
			Content:   content,
			ToolCalls: choice.Message.ToolCalls,
		}
		messages = append(messages, assistantMsg)

		for _, tc := range choice.Message.ToolCalls {
			toolCallData := map[string]interface{}{
				"tool_call_id": tc.ID,
				"name":         tc.Function.Name,
				"args":         tc.Function.Arguments,
			}
			writeSSE(map[string]interface{}{"tool_call": toolCallData})

			argsMap := make(map[string]interface{})
			json.Unmarshal([]byte(tc.Function.Arguments), &argsMap)

			tcInfo := ToolCallInfo{
				ID:   tc.ID,
				Name: tc.Function.Name,
				Args: argsMap,
			}

			call := agent.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: json.RawMessage(tc.Function.Arguments),
			}

			result, execErr := s.agentRegistry.Execute(ctx, call)
			if execErr != nil {
				tcInfo.Error = execErr.Error()
				writeSSE(map[string]interface{}{"tool_result": tcInfo})
			} else {
				tcInfo.Result = &toolResponseData{
					Content: result.Content,
					IsError: result.IsError,
					Meta:    result.Meta,
				}
				writeSSE(map[string]interface{}{"tool_result": tcInfo})
			}

			toolCalls = append(toolCalls, tcInfo)

			messages = append(messages, orchestrator.Message{
				Role:       "tool",
				Content:    result.Content,
				ToolCallID: tc.ID,
				Name:       tc.Function.Name,
			})
		}

		finalContent = ""
	}

	if finalContent == "" && len(toolCalls) > 0 {
		finalContent = "(opérations terminées)"
	}

	writeJSONResp, _ := json.Marshal(ShellChatResponse{
		Content:   finalContent,
		ToolCalls: toolCalls,
	})
	writeSSE(map[string]interface{}{"done": true, "response": string(writeJSONResp)})
}

func (s *Server) handleShellChatNonStream(w http.ResponseWriter, orb *orchestrator.Orchestrator, req ShellChatRequest) {
	ctx := context.Background()
	messages := []orchestrator.Message{
		{Role: "user", Content: req.Message},
	}

	var finalContent string
	var toolCalls []ToolCallInfo

	for i := 0; i < maxShellToolIterations; i++ {
		resp, err := orb.SendWithTools(messages)
		if err != nil {
			writeError(w, err.Error(), http.StatusInternalServerError)
			return
		}

		choice := resp.Choices[0]
		content := cleanThinkingTags(choice.Message.Content)

		if content != "" {
			finalContent = content
		}

		if len(choice.Message.ToolCalls) == 0 {
			break
		}

		assistantMsg := orchestrator.Message{
			Role:      "assistant",
			Content:   content,
			ToolCalls: choice.Message.ToolCalls,
		}
		messages = append(messages, assistantMsg)

		for _, tc := range choice.Message.ToolCalls {
			argsMap := make(map[string]interface{})
			json.Unmarshal([]byte(tc.Function.Arguments), &argsMap)

			tcInfo := ToolCallInfo{
				ID:   tc.ID,
				Name: tc.Function.Name,
				Args: argsMap,
			}

			call := agent.ToolCall{
				ID:        tc.ID,
				Name:      tc.Function.Name,
				Arguments: json.RawMessage(tc.Function.Arguments),
			}

			result, execErr := s.agentRegistry.Execute(ctx, call)
			if execErr != nil {
				tcInfo.Error = execErr.Error()
			} else {
				tcInfo.Result = &toolResponseData{
					Content: result.Content,
					IsError: result.IsError,
					Meta:    result.Meta,
				}
			}

			toolCalls = append(toolCalls, tcInfo)

			messages = append(messages, orchestrator.Message{
				Role:       "tool",
				Content:    result.Content,
				ToolCallID: tc.ID,
				Name:       tc.Function.Name,
			})
		}

		finalContent = ""
	}

	if finalContent == "" && len(toolCalls) > 0 {
		finalContent = "(tool calls completed, no text response)"
	}

	writeJSON(w, ShellChatResponse{
		Content:   finalContent,
		ToolCalls: toolCalls,
	})
}