fix: only set automatic close on events added to session and emitted … (#343)

fix: only set automatic close on events added to session and emitted by flow agent

Change-Id: I48c12915b0f2ee5606262bbf0857eb1a571faf10

Author: shentongmartin

adk/chatmodel.go

@@ -302,7 +302,6 @@ func (h *cbHandler) onChatModelEndWithStreamOutput(ctx context.Context,
 	}
 	out := schema.StreamReaderWithConvert(output, cvt)
 	event := EventFromMessage(nil, out, schema.Assistant, "")
-	setAutomaticClose(event)
 	h.Send(event)
 
 	return ctx
@@ -332,7 +331,6 @@ func (h *cbHandler) onToolEndWithStreamOutput(ctx context.Context,
 	}
 	out := schema.StreamReaderWithConvert(output, cvt)
 	event := EventFromMessage(nil, out, schema.Tool, runInfo.Name)
-	setAutomaticClose(event)
 	h.Send(event)
 
 	return ctx
@@ -441,7 +439,6 @@ func (a *ChatModelAgent) buildRunFunc(ctx context.Context) runFunc {
 							// copy the stream first because when setting output to session, the stream will be consumed
 							ss := msgStream.Copy(2)
 							event = EventFromMessage(msg, ss[1], schema.Assistant, "")
-							setAutomaticClose(event)
 							msgStream = ss[0]
 						} else {
 							event = EventFromMessage(msg, nil, schema.Assistant, "")

adk/flow.go

@@ -298,6 +298,7 @@ func (a *flowAgent) Run(ctx context.Context, input *AgentInput, opts ...AgentRun
 			// copy the event so that the copied event's stream is exclusive for any potential consumer
 			// copy before adding to session because once added to session it's stream could be consumed by genAgentInput at any time
 			copied := copyAgentEvent(event)
+			setAutomaticClose(event)
 			runCtx.session.addEvent(event)
 			setAutomaticClose(copied)
 			generator.Send(copied)