Agent-System-Backend/internal/runtime/squad_runner.go

package runtime

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	agentllm "agent-native/runtime/llm"
	"agent-squad/squad"
	"agent-system-backend/internal/squadadapter"
	core "agent-system-core"
)

// SquadRunner 是 runtime.entry=squad 时的对话入口。
// 小队先分诊：简单任务由 assistant 单步完成，复杂任务交给 leader。步骤都跑在内核上。
type SquadRunner struct {
	Stream  agentllm.StreamFunc
	Router  squad.Router
	Checker squad.Checker
	Hitl    squad.HitlPolicy
}

func (s *SquadRunner) Name() string { return "squad" }

func (s *SquadRunner) Execute(rc RunCtx) error {
	if rc.Context().Err() != nil {
		return rc.Context().Err()
	}
	prof := rc.Profile()
	maxTurns, maxTools := loopCaps(prof.Budget)
	produced := &artifactList{}
	var streamed strings.Builder
	var toolParts []core.Part
	textIdx := 0
	onEvent := func(ev squad.ProcessEvent) {
		switch ev.Type {
		case "text_delta":
			if ev.Text == "" {
				return
			}
			streamed.WriteString(ev.Text)
			_ = rc.Emit(core.EvMessageDelta, core.MessageDeltaPayload{Index: textIdx, Text: ev.Text})
		case "reasoning_delta":
			if ev.Text != "" {
				_ = rc.Emit(core.EvThinkingDelta, core.ThinkingDeltaPayload{Text: ev.Text})
			}
		case "tool_execution_start":
			_ = rc.Emit(core.EvToolCall, core.ToolCallPayload{CallID: ev.Tool, Name: ev.Tool, Arguments: ev.Detail})
			toolParts = append(toolParts, core.Part{Type: core.PartToolCall, Process: true, ToolCall: &core.ToolCallPart{ID: ev.Tool, Name: ev.Tool, Arguments: ev.Detail}})
		case "tool_execution_end":
			_ = rc.Emit(core.EvToolResult, core.ToolResultPayload{CallID: ev.Tool, Status: "ok", Summary: truncateRun(ev.Detail, 160)})
			toolParts = append(toolParts, core.Part{Type: core.PartToolResult, Process: true, ToolResult: &core.ToolResultPart{CallID: ev.Tool, Status: "ok", Summary: truncateRun(ev.Detail, 160)}})
		case "usage":
			if sink, ok := rc.(interface{ AddUsage(core.Usage) }); ok {
				sink.AddUsage(core.Usage{
					PromptTokens: ev.PromptTokens, CompletionTokens: ev.CompletionTokens, TotalTokens: ev.PromptTokens + ev.CompletionTokens,
					LastPromptTokens: ev.LastPromptTokens, LastCompletionTokens: ev.LastCompletionTokens,
					CacheReadTokens: ev.CacheReadTokens, CacheWriteTokens: ev.CacheWriteTokens,
				})
			}
		}
	}
	router := s.Router
	if router == nil {
		router = newModelRouter(rc)
	}
	res, err := squadadapter.RunEntry(rc, squadadapter.EntryOptions{
		DataDir:      rc.Settings().GetString("storage.dataDir", "./data"),
		Goal:         entryGoal(rc),
		HostTools:    bridgeHostTools(rc, produced),
		OnEvent:      onEvent,
		Stream:       s.Stream,
		Router:       router,
		Checker:      s.Checker,
		Hitl:         s.Hitl,
		MaxTurns:     maxTurns,
		MaxToolCalls: maxTools,
		Model:        ResolveChatModel(prof, rc.Settings()).Model,
	})
	if rc.Context().Err() != nil {
		return rc.Context().Err()
	}
	if err != nil {
		return err
	}
	reply := strings.TrimSpace(streamed.String())
	if reply == "" {
		reply = strings.TrimSpace(res.Summary)
	}
	if reply == "" {
		reply = "已完成。"
	}
	if claimsFalseDeliverable(reply, produced.items, rc) {
		return fmt.Errorf("声称已生成文件，但没有登记对应工件")
	}
	if strings.TrimSpace(res.Body) != "" {
		if _, werr := rc.WriteArtifact(res.Name, "document", "text/markdown", []byte(res.Body)); werr != nil {
			return werr
		}
	}
	parts := append([]core.Part{}, toolParts...)
	parts = append(parts, core.Part{Type: core.PartText, Text: reply})
	for _, ref := range produced.items {
		parts = append(parts, core.Part{Type: core.PartArtifactRef, Artifact: ref})
	}
	return rc.SetAssistantParts(parts)
}

func entryGoal(rc RunCtx) string {
	var b strings.Builder
	b.WriteString(currentUserText(rc))
	if listed, ok := rc.(interface{ WorkspaceInputs() []string }); ok {
		if paths := listed.WorkspaceInputs(); len(paths) > 0 {
			b.WriteString("\n工作区 inputs：")
			b.WriteString(strings.Join(paths, "、"))
		}
	}
	names := make([]string, 0, len(rc.Attachments()))
	for _, a := range rc.Attachments() {
		if a.Meta.Filename != "" {
			names = append(names, a.Meta.Filename)
		}
	}
	if len(names) > 0 {
		b.WriteString("\n附件：")
		b.WriteString(strings.Join(names, "、"))
	}
	return b.String()
}

type modelRouter struct {
	stream agentllm.StreamFunc
	model  string
}

func newModelRouter(rc RunCtx) modelRouter {
	route := ResolveLightweightModel(rc.Profile(), rc.Settings())
	if route.Model == "" || route.Base == "" {
		return modelRouter{}
	}
	client := openLLMClient(route.Base, route.Key)
	return modelRouter{stream: client.Stream, model: route.Model}
}

func (r modelRouter) Route(ctx context.Context, goal squad.Goal) (*squad.RouteDecision, error) {
	if r.stream == nil || r.model == "" {
		return &squad.RouteDecision{Mode: squad.RoutePlan, Family: "none", Reason: "未配置分诊模型，按复杂任务处理"}, nil
	}
	prompt := `只输出一个 JSON 对象，不要解释。字段 mode 取 direct 或 plan，family 取 minutes、report、proposal 或 none，reason 用一句话。简单问答和单步改写用 direct。需要成文、多步调研或会议材料用 plan。目标：` + goal.Text
	var text string
	for ev := range r.stream(ctx, r.model, []agentllm.Message{{Role: agentllm.RoleUser, Content: prompt}}, nil, agentllm.StreamOptions{}) {
		switch ev.Kind {
		case agentllm.KindTextDelta:
			text += ev.Text
		case agentllm.KindMessageEnd:
			if ev.Message.Content != "" {
				text = ev.Message.Content
			}
		case agentllm.KindError:
			return &squad.RouteDecision{Mode: squad.RoutePlan, Family: "none", Reason: "分诊调用失败，按复杂任务处理"}, nil
		}
	}
	return parseRoute(text), nil
}

func parseRoute(text string) *squad.RouteDecision {
	dec := &squad.RouteDecision{Mode: squad.RoutePlan, Family: "none", Reason: "分诊结果无法解析，按复杂任务处理"}
	start := strings.Index(text, "{")
	end := strings.LastIndex(text, "}")
	if start < 0 || end <= start {
		return dec
	}
	var raw struct {
		Mode   string `json:"mode"`
		Family string `json:"family"`
		Reason string `json:"reason"`
	}
	if json.Unmarshal([]byte(text[start:end+1]), &raw) != nil {
		return dec
	}
	switch raw.Mode {
	case string(squad.RouteDirect):
		dec.Mode = squad.RouteDirect
	case string(squad.RoutePlan):
		dec.Mode = squad.RoutePlan
	default:
		return dec
	}
	if raw.Family != "" {
		dec.Family = raw.Family
	}
	if raw.Reason != "" {
		dec.Reason = raw.Reason
	}
	return dec
}