fix(adk): stabilize session message ids

Change-Id: I898c2fe90d3457176f59f3641b6562f36f16b282

Author: shentongmartin

adk/chatmodel.go

@@ -280,6 +280,12 @@ func newDefaultGenModelInput[M MessageType]() TypedGenModelInput[M] {
 	}
 }
 
+func ensureGeneratedMessageIDs[M MessageType](messages []M) {
+	for _, msg := range messages {
+		EnsureMessageID(msg)
+	}
+}
+
 // TypedChatModelAgentState represents the state of a chat model agent during conversation.
 // This is the primary state type for both TypedChatModelAgentMiddleware and AgentMiddleware callbacks.
 type TypedChatModelAgentState[M MessageType] struct {
@@ -1121,6 +1127,9 @@ func (a *TypedChatModelAgent[M]) buildNoToolsRunFunc(_ context.Context) (typedRu
 			if err != nil {
 				return nil, err
 			}
+			if p.sessionEvents {
+				ensureGeneratedMessageIDs(messages)
+			}
 			if err := compose.ProcessState(ctx, func(_ context.Context, st *typedState[M]) error {
 				st.Messages = append(st.Messages, messages...)
 				return nil
@@ -1289,6 +1298,9 @@ func (a *TypedChatModelAgent[M]) buildMessageReActRunFunc(_ context.Context, bc
 					if genErr != nil {
 						return nil, genErr
 					}
+					if mp.sessionEvents {
+						ensureGeneratedMessageIDs(messages)
+					}
 					return &reactInput{
 						Messages: messages,
 					}, nil
@@ -1443,6 +1455,9 @@ func (a *TypedChatModelAgent[M]) buildAgenticReActRunFunc(_ context.Context, bc
 					if genErr != nil {
 						return nil, genErr
 					}
+					if ap.sessionEvents {
+						ensureGeneratedMessageIDs(messages)
+					}
 					return &agenticReactInput{
 						Messages: messages,
 					}, nil

adk/middlewares/patchtoolcalls/patchtoolcalls.go

@@ -165,6 +165,8 @@ type normalizationPlan[M adk.MessageType] struct {
 }
 
 func buildMessageNormalizationPlan(ctx context.Context, cfg Config, messages []*schema.Message) (*normalizationPlan[*schema.Message], error) {
+	ensureMessageIDs(messages)
+
 	counts := analyzeMessages(messages)
 	if cfg.Strict && counts.hasMismatch() {
 		return nil, counts.strictError()
@@ -248,6 +250,12 @@ func analyzeMessages(messages []*schema.Message) mismatchCounts {
 	return counts
 }
 
+func ensureMessageIDs[M adk.MessageType](messages []M) {
+	for _, msg := range messages {
+		adk.EnsureMessageID(msg)
+	}
+}
+
 func keptMessages(messages []*schema.Message, cfg Config) []bool {
 	keep := make([]bool, len(messages))
 	previousCalls := make(map[string]struct{})
@@ -281,6 +289,8 @@ func keptMessages(messages []*schema.Message, cfg Config) []bool {
 }
 
 func buildAgenticNormalizationPlan(ctx context.Context, cfg Config, messages []*schema.AgenticMessage) (*normalizationPlan[*schema.AgenticMessage], error) {
+	ensureMessageIDs(messages)
+
 	counts := analyzeAgenticMessages(messages)
 	if cfg.Strict && counts.hasMismatch() {
 		return nil, counts.strictError()

adk/runner.go

@@ -319,11 +319,13 @@ func prepareRunnerSessionRun[M MessageType]( //nolint:revive // argument-limit
 		TurnID:    state.turnID,
 		Lifecycle: &LifecycleEvent{Scope: LifecycleScopeSession, State: SessionRunStateRunning},
 	}
-	if err := assignSessionEventID(ctx, runningEvent, state.sessionConfig.EventIDGenerator); err != nil {
+	err = assignSessionEventID(ctx, runningEvent, state.sessionConfig.EventIDGenerator)
+	if err != nil {
 		_ = state.sessionHandle.close(ctx)
 		return nil, err
 	}
-	if err := appendRunnerSessionControlEvent(ctx, state, runningEvent, ""); err != nil {
+	err = appendRunnerSessionControlEvent(ctx, state, runningEvent, "")
+	if err != nil {
 		_ = state.sessionHandle.close(ctx)
 		return nil, err
 	}
@@ -352,7 +354,7 @@ func prepareRunnerSessionRun[M MessageType]( //nolint:revive // argument-limit
 	return state, nil
 }
 
-func prepareRunnerSessionResume[M MessageType](
+func prepareRunnerSessionResume[M MessageType]( //nolint:revive // argument-limit
 	ctx context.Context,
 	checkPointStore CheckPointStore,
 	sessionID string,
@@ -568,12 +570,12 @@ func typedRunnerRunImpl[M MessageType](a TypedAgent[M], enableStreaming bool, st
 		// Capture caller-provided messages BEFORE prepending history. These will be
 		// emitted as session events at turn start so they appear in the event log.
 		sessionState.inputMessages = append([]M{}, messages...)
-		// Assign eino message IDs to input messages (needed for BeforeMessageID references
-		// emitted by middlewares that anchor on user messages).
-		for _, msg := range sessionState.inputMessages {
+		messages = append(append([]M{}, sessionState.latestState.Messages...), sessionState.inputMessages...)
+		// Assign IDs before messages can be both persisted and inspected by
+		// middleware, avoiding concurrent lazy ID mutation during event snapshotting.
+		for _, msg := range messages {
 			EnsureMessageID(msg)
 		}
-		messages = append(append([]M{}, sessionState.latestState.Messages...), sessionState.inputMessages...)
 		o.sessionValues = mergeSessionValues(sessionState.latestState.SessionValues, o.sessionValues)
 		opts = append(opts, withEnableSessionEvents())
 		opts = append(opts, withEnableInternalTimelineEvents())

adk/wrappers.go

@@ -911,9 +911,13 @@ func GetMessageID[M MessageType](msg M) string {
 func EnsureMessageID[M MessageType](msg M) {
 	switch v := any(msg).(type) {
 	case *schema.Message:
-		v.Extra = internal.EnsureMessageID(v.Extra)
+		if internal.GetMessageID(v.Extra) == "" {
+			v.Extra = internal.EnsureMessageID(v.Extra)
+		}
 	case *schema.AgenticMessage:
-		v.Extra = internal.EnsureMessageID(v.Extra)
+		if internal.GetMessageID(v.Extra) == "" {
+			v.Extra = internal.EnsureMessageID(v.Extra)
+		}
 	}
 }