MuyueWorkspace/internal/api/handlers_shell_chat.go

package api

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

	"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 toString(v interface{}) string {
	if v == nil {
		return ""
	}
	s, _ := v.(string)
	return s
}

func toBool(v interface{}) bool {
	if v == nil {
		return false
	}
	b, _ := v.(bool)
	return b
}

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) {
	SetupSSEHeaders(w)
	flusher, canFlush := w.(http.Flusher)
	sseWriter := NewSSEWriter(w)

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

	engine := NewChatEngine(orb, s.agentRegistry, s.agentToolsJSON)

	var toolCalls []ToolCallInfo
	engine.OnChunk(func(data map[string]interface{}) {
		if data == nil {
			return
		}
		sseWriter.Write(data)
		if canFlush {
			flusher.Flush()
		}
		if tc, ok := data["tool_call"].(map[string]interface{}); ok {
			argsMap := make(map[string]interface{})
			if args, ok := tc["args"].(string); ok {
				json.Unmarshal([]byte(args), &argsMap)
			}
			toolCalls = append(toolCalls, ToolCallInfo{
				ID:   toString(tc["tool_call_id"]),
				Name: toString(tc["name"]),
				Args: argsMap,
			})
		}
		if tr, ok := data["tool_result"].(map[string]interface{}); ok {
			tcID := toString(tr["tool_call_id"])
			for i := range toolCalls {
				if toolCalls[i].ID == tcID {
					if err, ok := tr["is_error"].(bool); ok && err {
						toolCalls[i].Error = toString(tr["content"])
					} else {
						toolCalls[i].Result = &toolResponseData{
							Content: toString(tr["content"]),
							IsError: toBool(tr["is_error"]),
						}
					}
					break
				}
			}
		}
	})

	finalContent, _, _, err := engine.RunWithTools(ctx, messages)
	if err != nil {
		sseWriter.Write(map[string]interface{}{"error": err.Error()})
		return
	}

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

	writeJSONResp, _ := json.Marshal(ShellChatResponse{
		Content:   finalContent,
		ToolCalls: toolCalls,
	})
	sseWriter.Write(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},
	}

	engine := NewChatEngine(orb, s.agentRegistry, s.agentToolsJSON)

	finalContent, err := engine.RunNonStream(ctx, messages)
	if err != nil {
		writeError(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if finalContent == "" {
		finalContent = "(tool calls completed, no text response)"
	}

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