Add files via upload

internal/multiagent/sub_agent_context.go

@@ -1,26 +1,81 @@
 package multiagent
 
 import (
+	"context"
+	"encoding/json"
 	"strings"
 
 	"cyberstrike-ai/internal/agent"
+
+	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/components/tool"
 )
 
 const defaultSubAgentUserContextMaxRunes = 2000
 
-// buildUserContextForSubAgent collects all user messages from conversation
-// history plus the current user message, and returns a formatted string to
-// append to sub-agent instructions. This ensures sub-agents always have
-// access to the full user intent (target URLs, scope, etc.) even when the
-// orchestrator forgets to include them in the task description.
+// taskContextEnrichMiddleware intercepts "task" tool calls on the orchestrator
+// and appends the user's original conversation messages to the task description.
+// This ensures sub-agents always receive the full user intent (target URLs,
+// scope, etc.) even when the orchestrator forgets to include them.
 //
-// maxRunes controls the character budget for the user-context body:
-//   - 0 uses defaultSubAgentUserContextMaxRunes
-//   - negative disables injection (returns "")
-//
-// When truncation is needed, the first and last user messages are each
-// allocated half the budget so neither is lost entirely.
-func buildUserContextForSubAgent(userMessage string, history []agent.ChatMessage, maxRunes int) string {
+// Design: user context is injected into the task description (per-task), NOT
+// into the sub-agent's Instruction (system prompt). This keeps sub-agent
+// Instructions clean as pure role definitions while attaching context to the
+// specific delegation — aligned with Claude Code's agent design philosophy.
+type taskContextEnrichMiddleware struct {
+	adk.BaseChatModelAgentMiddleware
+	supplement string // pre-built user context block
+}
+
+// newTaskContextEnrichMiddleware returns a middleware that enriches task
+// descriptions with user conversation context. Returns nil if disabled
+// (maxRunes < 0) or no user messages exist.
+func newTaskContextEnrichMiddleware(userMessage string, history []agent.ChatMessage, maxRunes int) adk.ChatModelAgentMiddleware {
+	supplement := buildUserContextSupplement(userMessage, history, maxRunes)
+	if supplement == "" {
+		return nil
+	}
+	return &taskContextEnrichMiddleware{supplement: supplement}
+}
+
+func (m *taskContextEnrichMiddleware) WrapInvokableToolCall(
+	ctx context.Context,
+	endpoint adk.InvokableToolCallEndpoint,
+	tCtx *adk.ToolContext,
+) (adk.InvokableToolCallEndpoint, error) {
+	if tCtx == nil || !strings.EqualFold(strings.TrimSpace(tCtx.Name), "task") {
+		return endpoint, nil
+	}
+	return func(ctx context.Context, argumentsInJSON string, opts ...tool.Option) (string, error) {
+		enriched := m.enrichTaskDescription(argumentsInJSON)
+		return endpoint(ctx, enriched, opts...)
+	}, nil
+}
+
+// enrichTaskDescription parses the task JSON arguments, appends user context
+// to the "description" field, and re-serializes. Falls back to the original
+// JSON if parsing fails or no description field exists.
+func (m *taskContextEnrichMiddleware) enrichTaskDescription(argsJSON string) string {
+	var raw map[string]interface{}
+	if err := json.Unmarshal([]byte(argsJSON), &raw); err != nil {
+		return argsJSON
+	}
+	desc, ok := raw["description"].(string)
+	if !ok {
+		return argsJSON
+	}
+	raw["description"] = desc + m.supplement
+	enriched, err := json.Marshal(raw)
+	if err != nil {
+		return argsJSON
+	}
+	return string(enriched)
+}
+
+// buildUserContextSupplement collects user messages from conversation history
+// and the current message, returning a formatted block to append to task
+// descriptions. Returns "" if disabled or no user messages exist.
+func buildUserContextSupplement(userMessage string, history []agent.ChatMessage, maxRunes int) string {
 	if maxRunes < 0 {
 		return ""
 	}
@@ -46,17 +101,16 @@ func buildUserContextForSubAgent(userMessage string, history []agent.ChatMessage
 	}
 
 	joined := strings.Join(userMsgs, "\n---\n")
-
 	if len([]rune(joined)) > maxRunes {
 		joined = truncateKeepFirstLast(userMsgs, maxRunes)
 	}
 
-	return "\n\n## 本次会话用户原始请求（自动注入，确保你了解完整上下文）\n" + joined
+	return "\n\n## 会话上下文（自动补充，确保你了解用户完整意图）\n" + joined
 }
 
 // truncateKeepFirstLast keeps the first and last user messages, giving each
 // half the rune budget. The first message typically contains target info;
-// the last is the current instruction.
+// the last contains the current instruction.
 func truncateKeepFirstLast(msgs []string, maxRunes int) string {
 	if len(msgs) == 1 {
 		return truncateRunes(msgs[0], maxRunes)