Add files via upload

internal/multiagent/eino_model_facing_trace.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 
 	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/schema"
 )
 
 // modelFacingTraceHolder 保存「即将送入 ChatModel」的消息快照（已走 summarization / reduction / orphan 修剪等），
@@ -43,6 +44,19 @@ func (h *modelFacingTraceHolder) storeFromState(state *adk.ChatModelAgentState)
 	h.mu.Unlock()
 }
 
+func (h *modelFacingTraceHolder) storeFromAgenticState(state *adk.TypedChatModelAgentState[*schema.AgenticMessage]) {
+	if h == nil || state == nil || len(state.Messages) == 0 {
+		return
+	}
+	cloned := cloneADKMessagesForTrace(AgenticMessagesToEino(state.Messages))
+	if len(cloned) == 0 {
+		return
+	}
+	h.mu.Lock()
+	h.msgs = cloned
+	h.mu.Unlock()
+}
+
 func cloneADKMessagesForTrace(msgs []adk.Message) []adk.Message {
 	if len(msgs) == 0 {
 		return nil
@@ -82,3 +96,29 @@ func (m *modelFacingTraceMiddleware) BeforeModelRewriteState(
 	}
 	return ctx, state, nil
 }
+
+type agenticModelFacingTraceMiddleware struct {
+	*adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]
+	holder *modelFacingTraceHolder
+}
+
+func newAgenticModelFacingTraceMiddleware(holder *modelFacingTraceHolder) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] {
+	if holder == nil {
+		return nil
+	}
+	return &agenticModelFacingTraceMiddleware{
+		TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{},
+		holder:                            holder,
+	}
+}
+
+func (m *agenticModelFacingTraceMiddleware) BeforeModelRewriteState(
+	ctx context.Context,
+	state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
+	mc *adk.TypedModelContext[*schema.AgenticMessage],
+) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) {
+	if m.holder != nil && state != nil {
+		m.holder.storeFromAgenticState(state)
+	}
+	return ctx, state, nil
+}