MuyueWorkspace/internal/api/handlers_ai_task.go

package api

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"runtime"
	"strings"
	"time"

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

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

	var body struct {
		Task string `json:"task"`
		Tool string `json:"tool,omitempty"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeError(w, err.Error(), http.StatusBadRequest)
		return
	}
	if body.Task == "" {
		writeError(w, "task is required", http.StatusBadRequest)
		return
	}

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

	orb.SetSystemPrompt(buildAITaskSystemPrompt())
	orb.SetTools(s.shellAgentToolsJSON)

	ctx, cancel := context.WithTimeout(context.Background(), 180*time.Second)
	defer cancel()

	messages := []orchestrator.Message{
		{Role: "user", Content: orchestrator.TextContent(buildAITaskPrompt(body.Task, body.Tool))},
	}

	engine := NewChatEngine(orb, s.shellAgentRegistry, s.shellAgentToolsJSON)
	finalContent, err := engine.RunNonStream(ctx, messages)
	if err != nil {
		writeError(w, "AI task failed: "+err.Error(), http.StatusInternalServerError)
		return
	}

	s.consumption.Record(engine.ProviderName(), engine.TotalTokens)

	parsed := parseAIJSONResponse(finalContent)

	writeJSON(w, map[string]interface{}{
		"status":  "ok",
		"raw":     finalContent,
		"result":  parsed,
		"tokens":  engine.TotalTokens,
	})
}

func buildAITaskSystemPrompt() string {
	return fmt.Sprintf(`You are a system administration assistant. You have access to a terminal tool to run commands on the host system.

IMPORTANT RULES:
- You MUST respond ONLY with valid JSON. No markdown, no code fences, no extra text.
- Always run the actual commands needed to complete the task.
- Be thorough: check versions, verify installations, compare with latest releases.

OS: %s/%s
Date: %s
`, runtime.GOOS, runtime.GOARCH, time.Now().Format("2006-01-02"))
}

func buildAITaskPrompt(task, tool string) string {
	switch task {
	case "check_tools":
		return `Check the following tools on this system. For each tool, determine:
1. Is it installed? Run "which <tool>" or "<tool> --version"
2. If installed, what is the current version?
3. What is the latest available version? Check GitHub releases API or official sources.

Tools to check: crush, claude, git, node, npm, pnpm, python3, pip3, uv, go, docker, gh, starship, npx

Run the commands needed, then respond with ONLY this JSON structure (no markdown fences):
{
  "tools": [
    {"name": "tool_name", "installed": true/false, "version": "x.y.z", "latest": "a.b.c", "needs_update": true/false, "category": "ai|runtime|vcs|devops|prompt"}
  ]
}`

	case "install_tool":
		return fmt.Sprintf(`Install the tool "%s" on this system.

Steps:
1. Check if it's already installed: run "which %s" and "%s --version"
2. If not installed, determine the best installation method for this OS
3. Run the installation command
4. Verify the installation succeeded

Respond with ONLY this JSON (no markdown fences):
{
  "tool": "%s",
  "installed": true/false,
  "version": "installed version or empty",
  "message": "what was done",
  "error": "error message or empty"
}`, tool, tool, tool, tool)

	case "update_tool":
		return fmt.Sprintf(`Update the tool "%s" to its latest version on this system.

Steps:
1. Check current version: run "%s --version"
2. Find the latest version available
3. Run the update/upgrade command
4. Verify the new version

Respond with ONLY this JSON (no markdown fences):
{
  "tool": "%s",
  "previous_version": "old version",
  "version": "new version",
  "updated": true/false,
  "message": "what was done",
  "error": "error message or empty"
}`, tool, tool, tool)

	default:
		return task
	}
}

func parseAIJSONResponse(content string) interface{} {
	cleaned := content

	if idx := strings.Index(cleaned, "```json"); idx != -1 {
		cleaned = cleaned[idx+7:]
		if end := strings.Index(cleaned, "```"); end != -1 {
			cleaned = cleaned[:end]
		}
	} else if idx := strings.Index(cleaned, "```"); idx != -1 {
		cleaned = cleaned[idx+3:]
		if end := strings.Index(cleaned, "```"); end != -1 {
			cleaned = cleaned[:end]
		}
	}

	cleaned = strings.TrimSpace(cleaned)

	jsonStart := strings.Index(cleaned, "{")
	jsonEnd := strings.LastIndex(cleaned, "}")
	if jsonStart != -1 && jsonEnd > jsonStart {
		cleaned = cleaned[jsonStart : jsonEnd+1]
	}

	var result interface{}
	if err := json.Unmarshal([]byte(cleaned), &result); err != nil {
		return map[string]interface{}{
			"raw":   content,
			"error": "failed to parse AI response as JSON",
		}
	}
	return result
}