Merge pull request #1349 from trheyi/main

Refactor message handling for improved streaming and event management
This commit is contained in:
Max 2025-11-26 18:31:17 +08:00 committed by GitHub
commit 3d60ffbaa2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 1072 additions and 302 deletions

View file

@ -704,7 +704,7 @@ func (ast *Assistant) sendAgentStreamStart(ctx *context.Context, handler message
} }
// Build the start data // Build the start data
startData := message.StreamStartData{ startData := message.EventStreamStartData{
ContextID: ctx.ID, ContextID: ctx.ID,
ChatID: ctx.ChatID, ChatID: ctx.ChatID,
TraceID: ctx.TraceID(), TraceID: ctx.TraceID(),
@ -731,7 +731,7 @@ func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler message.S
return return
} }
endData := &message.StreamEndData{ endData := &message.EventStreamEndData{
RequestID: ctx.RequestID(), RequestID: ctx.RequestID(),
ContextID: ctx.ID, ContextID: ctx.ID,
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),

View file

@ -82,7 +82,7 @@ type streamState struct {
func (s *streamState) handleStreamStart(data []byte) int { func (s *streamState) handleStreamStart(data []byte) int {
// Send event message to indicate stream has started // Send event message to indicate stream has started
// This is a lifecycle event, CUI clients can show it, OpenAI clients will ignore it // This is a lifecycle event, CUI clients can show it, OpenAI clients will ignore it
var startData message.StreamStartData var startData message.EventStreamStartData
err := jsoniter.Unmarshal(data, &startData) err := jsoniter.Unmarshal(data, &startData)
if err != nil { if err != nil {
log.Error("Failed to unmarshal stream start data: %v", err) log.Error("Failed to unmarshal stream start data: %v", err)
@ -95,17 +95,17 @@ func (s *streamState) handleStreamStart(data []byte) int {
// handleGroupStart handles group start event // handleGroupStart handles group start event
func (s *streamState) handleGroupStart(data []byte) int { func (s *streamState) handleGroupStart(data []byte) int {
// Parse group start data first to get the group ID // Parse group start data first to get the group ID
var startData message.GroupStartData var startData message.EventMessageStartData
if err := jsoniter.Unmarshal(data, &startData); err != nil { if err := jsoniter.Unmarshal(data, &startData); err != nil {
log.Error("Failed to unmarshal group start data: %v", err) log.Error("Failed to unmarshal group start data: %v", err)
return 0 return 0
} }
// Use the group ID from the start data, or generate one if not provided // Use the message ID from the start data, or generate one if not provided
groupID := startData.GroupID groupID := startData.MessageID
if groupID == "" { if groupID == "" {
groupID = generateMessageID() groupID = generateMessageID()
startData.GroupID = groupID startData.MessageID = groupID
} }
// Initialize group state with the correct group ID // Initialize group state with the correct group ID
@ -138,11 +138,11 @@ func (s *streamState) handleText(data []byte) int {
s.messageSeq++ s.messageSeq++
// Send delta message // Send delta message
// - ID: Sequential message ID (e.g., msg_001, msg_002) for readability // - ChunkID: Sequential chunk ID (C1, C2, C3...) for this fragment
// - GroupID: Same for all chunks of this logical message (frontend merges by group_id) // - MessageID: Same for all chunks of this logical message (frontend merges by message_id)
msg := &message.Message{ msg := &message.Message{
ID: s.generateSequentialID(), // Sequential ID for this chunk ChunkID: s.generateSequentialID(), // Sequential ID for this chunk
GroupID: s.currentGroupID, // Group ID for merging (all chunks share this) MessageID: s.currentGroupID, // Message ID for merging (all chunks share this)
Type: message.TypeText, Type: message.TypeText,
Delta: true, Delta: true,
Props: map[string]interface{}{ Props: map[string]interface{}{
@ -173,11 +173,11 @@ func (s *streamState) handleThinking(data []byte) int {
s.messageSeq++ s.messageSeq++
// Send delta message // Send delta message
// - ID: Sequential message ID (e.g., msg_001, msg_002) for readability // - ChunkID: Sequential chunk ID (C1, C2, C3...) for this fragment
// - GroupID: Same for all chunks of this logical message (frontend merges by group_id) // - MessageID: Same for all chunks of this logical message (frontend merges by message_id)
msg := &message.Message{ msg := &message.Message{
ID: s.generateSequentialID(), // Sequential ID for this chunk ChunkID: s.generateSequentialID(), // Sequential ID for this chunk
GroupID: s.currentGroupID, // Group ID for merging (all chunks share this) MessageID: s.currentGroupID, // Message ID for merging (all chunks share this)
Type: message.TypeThinking, Type: message.TypeThinking,
Delta: true, Delta: true,
Props: map[string]interface{}{ Props: map[string]interface{}{
@ -197,7 +197,7 @@ func (s *streamState) handleToolCall(data []byte) int {
// Tool calls are usually complete JSON objects // Tool calls are usually complete JSON objects
// Parse and send as tool_call message // Parse and send as tool_call message
msg := &message.Message{ msg := &message.Message{
ID: generateMessageID(), MessageID: generateMessageID(),
Type: message.TypeToolCall, Type: message.TypeToolCall,
Delta: true, Delta: true,
Props: map[string]interface{}{ Props: map[string]interface{}{
@ -241,9 +241,9 @@ func (s *streamState) handleGroupEnd(data []byte) int {
msgType = message.TypeText // Fallback to text if type not set msgType = message.TypeText // Fallback to text if type not set
} }
// Build GroupEndData with complete content // Build EventMessageEndData with complete content
endData := message.GroupEndData{ endData := message.EventMessageEndData{
GroupID: s.currentGroupID, // Use the group ID, not message ID MessageID: s.currentGroupID, // Use the message ID
Type: msgType, Type: msgType,
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),
DurationMs: durationMs, DurationMs: durationMs,
@ -271,7 +271,7 @@ func (s *streamState) handleGroupEnd(data []byte) int {
// handleStreamEnd handles stream end event // handleStreamEnd handles stream end event
func (s *streamState) handleStreamEnd(data []byte) int { func (s *streamState) handleStreamEnd(data []byte) int {
// Parse the stream end data // Parse the stream end data
var endData message.StreamEndData var endData message.EventStreamEndData
if err := jsoniter.Unmarshal(data, &endData); err != nil { if err := jsoniter.Unmarshal(data, &endData); err != nil {
log.Error("Failed to parse stream_end data: %v", err) log.Error("Failed to parse stream_end data: %v", err)
s.ctx.Flush() s.ctx.Flush()

View file

@ -9,6 +9,7 @@ import (
jsoniter "github.com/json-iterator/go" jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/plan" "github.com/yaoapp/gou/plan"
"github.com/yaoapp/kun/log" "github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/output/message"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/trace" "github.com/yaoapp/yao/trace"
@ -33,6 +34,7 @@ func New(parent context.Context, authorized *types.AuthorizedInfo, chatID, paylo
ID: generateContextID(), // Generate unique ID for the context ID: generateContextID(), // Generate unique ID for the context
Space: plan.NewMemorySharedSpace(), Space: plan.NewMemorySharedSpace(),
ChatID: chatID, ChatID: chatID,
IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context
} }
if payload == "" { if payload == "" {

View file

@ -185,9 +185,13 @@ func (ctx *Context) sendMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return bridge.JsException(v8ctx, "invalid message: "+err.Error()) return bridge.JsException(v8ctx, "invalid message: "+err.Error())
} }
// Generate unique ID if not provided // Generate unique MessageID if not provided
if msg.ID == "" { if msg.MessageID == "" {
msg.ID = output.GenerateID() if ctx.IDGenerator != nil {
msg.MessageID = ctx.IDGenerator.GenerateMessageID()
} else {
msg.MessageID = output.GenerateID()
}
} }
// Call ctx.Send // Call ctx.Send
@ -222,18 +226,22 @@ func (ctx *Context) sendGroupMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return bridge.JsException(v8ctx, "invalid group: "+err.Error()) return bridge.JsException(v8ctx, "invalid group: "+err.Error())
} }
// Generate group ID if not provided // Generate block ID if not provided
if group.ID == "" { if group.ID == "" {
if ctx.IDGenerator != nil {
group.ID = ctx.IDGenerator.GenerateBlockID()
} else {
group.ID = output.GenerateID() group.ID = output.GenerateID()
} }
}
// Send group_start event // Send group_start event
startTime := time.Now() startTime := time.Now()
startEvent := output.NewEventMessage( startEvent := output.NewEventMessage(
message.EventGroupStart, message.EventGroupStart,
"Group started", "Group started",
message.GroupStartData{ message.EventMessageStartData{
GroupID: group.ID, MessageID: group.ID,
Type: "mixed", // Mixed types in group Type: "mixed", // Mixed types in group
Timestamp: startTime.UnixMilli(), Timestamp: startTime.UnixMilli(),
}, },
@ -245,13 +253,17 @@ func (ctx *Context) sendGroupMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
return bridge.JsException(v8ctx, "Flush failed after group_start: "+err.Error()) return bridge.JsException(v8ctx, "Flush failed after group_start: "+err.Error())
} }
// Generate IDs for messages and set group_id // Generate MessageIDs for messages and set BlockID
for _, msg := range group.Messages { for _, msg := range group.Messages {
if msg.ID == "" { if msg.MessageID == "" {
msg.ID = output.GenerateID() if ctx.IDGenerator != nil {
msg.MessageID = ctx.IDGenerator.GenerateMessageID()
} else {
msg.MessageID = output.GenerateID()
} }
if msg.GroupID == "" { }
msg.GroupID = group.ID if msg.BlockID == "" {
msg.BlockID = group.ID
} }
} }
@ -267,8 +279,8 @@ func (ctx *Context) sendGroupMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
endEvent := output.NewEventMessage( endEvent := output.NewEventMessage(
message.EventGroupEnd, message.EventGroupEnd,
"Group completed", "Group completed",
message.GroupEndData{ message.EventMessageEndData{
GroupID: group.ID, MessageID: group.ID,
Type: "mixed", Type: "mixed",
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),
DurationMs: time.Since(startTime).Milliseconds(), DurationMs: time.Since(startTime).Milliseconds(),
@ -303,20 +315,24 @@ func (ctx *Context) sendGroupStartMethod(iso *v8go.Isolate) *v8go.FunctionTempla
groupType = args[0].String() groupType = args[0].String()
} }
// Get or generate group ID // Get or generate block ID
var groupID string var groupID string
if len(args) > 1 && args[1].IsString() { if len(args) > 1 && args[1].IsString() {
groupID = args[1].String() groupID = args[1].String()
} else {
if ctx.IDGenerator != nil {
groupID = ctx.IDGenerator.GenerateBlockID()
} else { } else {
groupID = output.GenerateID() groupID = output.GenerateID()
} }
}
// Send group_start event // Send group_start event
startEvent := output.NewEventMessage( startEvent := output.NewEventMessage(
message.EventGroupStart, message.EventGroupStart,
"Group started", "Group started",
message.GroupStartData{ message.EventMessageStartData{
GroupID: groupID, MessageID: groupID,
Type: groupType, Type: groupType,
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),
}, },
@ -361,8 +377,8 @@ func (ctx *Context) sendGroupEndMethod(iso *v8go.Isolate) *v8go.FunctionTemplate
endEvent := output.NewEventMessage( endEvent := output.NewEventMessage(
message.EventGroupEnd, message.EventGroupEnd,
"Group completed", "Group completed",
message.GroupEndData{ message.EventMessageEndData{
GroupID: groupID, MessageID: groupID,
Type: "mixed", Type: "mixed",
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),
DurationMs: 0, // Duration not tracked at this level DurationMs: 0, // Duration not tracked at this level

View file

@ -51,10 +51,21 @@ func parseMessage(v8ctx *v8go.Context, jsValue *v8go.Value) (*message.Message, e
msg.Props = props msg.Props = props
} }
// Optional fields // Optional fields - Streaming control
if id, ok := msgMap["id"].(string); ok { if chunkID, ok := msgMap["chunk_id"].(string); ok {
msg.ID = id msg.ChunkID = chunkID
} }
if messageID, ok := msgMap["message_id"].(string); ok {
msg.MessageID = messageID
}
if blockID, ok := msgMap["block_id"].(string); ok {
msg.BlockID = blockID
}
if threadID, ok := msgMap["thread_id"].(string); ok {
msg.ThreadID = threadID
}
// Delta control
if delta, ok := msgMap["delta"].(bool); ok { if delta, ok := msgMap["delta"].(bool); ok {
msg.Delta = delta msg.Delta = delta
} }
@ -67,9 +78,6 @@ func parseMessage(v8ctx *v8go.Context, jsValue *v8go.Value) (*message.Message, e
if typeChange, ok := msgMap["type_change"].(bool); ok { if typeChange, ok := msgMap["type_change"].(bool); ok {
msg.TypeChange = typeChange msg.TypeChange = typeChange
} }
if groupID, ok := msgMap["group_id"].(string); ok {
msg.GroupID = groupID
}
// Metadata (optional) // Metadata (optional)
if metadataMap, ok := msgMap["metadata"].(map[string]interface{}); ok { if metadataMap, ok := msgMap["metadata"].(map[string]interface{}); ok {
@ -142,10 +150,21 @@ func parseGroup(v8ctx *v8go.Context, jsValue *v8go.Value) (*message.Group, error
msg.Props = props msg.Props = props
} }
// Optional fields // Optional fields - Streaming control
if id, ok := msgMap["id"].(string); ok { if chunkID, ok := msgMap["chunk_id"].(string); ok {
msg.ID = id msg.ChunkID = chunkID
} }
if messageID, ok := msgMap["message_id"].(string); ok {
msg.MessageID = messageID
}
if blockID, ok := msgMap["block_id"].(string); ok {
msg.BlockID = blockID
}
if threadID, ok := msgMap["thread_id"].(string); ok {
msg.ThreadID = threadID
}
// Delta control
if delta, ok := msgMap["delta"].(bool); ok { if delta, ok := msgMap["delta"].(bool); ok {
msg.Delta = delta msg.Delta = delta
} }
@ -158,9 +177,6 @@ func parseGroup(v8ctx *v8go.Context, jsValue *v8go.Value) (*message.Group, error
if typeChange, ok := msgMap["type_change"].(bool); ok { if typeChange, ok := msgMap["type_change"].(bool); ok {
msg.TypeChange = typeChange msg.TypeChange = typeChange
} }
if groupID, ok := msgMap["group_id"].(string); ok {
msg.GroupID = groupID
}
// Metadata (optional) // Metadata (optional)
if metadataMap, ok := msgMap["metadata"].(map[string]interface{}); ok { if metadataMap, ok := msgMap["metadata"].(map[string]interface{}); ok {

View file

@ -7,6 +7,7 @@ import (
"github.com/yaoapp/gou/plan" "github.com/yaoapp/gou/plan"
"github.com/yaoapp/gou/store" "github.com/yaoapp/gou/store"
"github.com/yaoapp/yao/agent/output" "github.com/yaoapp/yao/agent/output"
"github.com/yaoapp/yao/agent/output/message"
"github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/openapi/oauth/types"
traceTypes "github.com/yaoapp/yao/trace/types" traceTypes "github.com/yaoapp/yao/trace/types"
) )
@ -207,6 +208,7 @@ type Context struct {
Skip *Skip `json:"skip,omitempty"` // Skip configuration (history, trace, etc.), nil means don't skip anything Skip *Skip `json:"skip,omitempty"` // Skip configuration (history, trace, etc.), nil means don't skip anything
trace traceTypes.Manager `json:"-"` // Trace manager, lazy initialized on first access trace traceTypes.Manager `json:"-"` // Trace manager, lazy initialized on first access
output *output.Output `json:"-"` // Output, it will be used to write response data to the client output *output.Output `json:"-"` // Output, it will be used to write response data to the client
IDGenerator *message.IDGenerator `json:"-"` // ID generator for this context (chunk, message, block, thread IDs)
// Model capabilities (set by assistant, used by output adapters) // Model capabilities (set by assistant, used by output adapters)
Capabilities *ModelCapabilities `json:"-"` // Model capabilities for the current connector Capabilities *ModelCapabilities `json:"-"` // Model capabilities for the current connector

View file

@ -18,25 +18,33 @@ import (
"github.com/yaoapp/yao/utils/jsonschema" "github.com/yaoapp/yao/utils/jsonschema"
) )
// startGroup starts a new group and sends group_start event // startMessage starts a new message and sends group_start event
func (gt *groupTracker) startGroup(groupType message.StreamChunkType, handler message.StreamFunc) { // Note: group_start/group_end events are used for backward compatibility
if gt.active { // but at LLM level they represent message boundaries, not Agent-level blocks
// End previous group first func (mt *messageTracker) startMessage(messageType message.StreamChunkType, handler message.StreamFunc) {
gt.endGroup(handler) if mt.active {
// End previous message first
mt.endMessage(handler)
} }
gt.active = true mt.active = true
gt.groupID = fmt.Sprintf("grp_%d", time.Now().UnixNano()) // Generate message ID using context's ID generator
gt.groupType = groupType if mt.idGenerator != nil {
gt.startTime = time.Now().UnixMilli() mt.messageID = mt.idGenerator.GenerateMessageID() // M1, M2, M3...
gt.chunkCount = 0 } else {
gt.toolCallInfo = nil // Fallback to global generator if no context generator
mt.messageID = message.GenerateNanoID()
}
mt.messageType = messageType
mt.startTime = time.Now().UnixMilli()
mt.chunkCount = 0
mt.toolCallInfo = nil
if handler != nil { if handler != nil {
startData := &message.GroupStartData{ startData := &message.EventMessageStartData{
GroupID: gt.groupID, MessageID: mt.messageID,
Type: string(groupType), Type: string(messageType),
Timestamp: gt.startTime, Timestamp: mt.startTime,
} }
if startJSON, err := jsoniter.Marshal(startData); err == nil { if startJSON, err := jsoniter.Marshal(startData); err == nil {
handler(message.ChunkGroupStart, startJSON) handler(message.ChunkGroupStart, startJSON)
@ -44,24 +52,30 @@ func (gt *groupTracker) startGroup(groupType message.StreamChunkType, handler me
} }
} }
// startToolCallGroup starts a new tool call group with tool call info // startToolCallMessage starts a new tool call message with tool call info
func (gt *groupTracker) startToolCallGroup(toolCallInfo *message.GroupToolCallInfo, handler message.StreamFunc) { func (mt *messageTracker) startToolCallMessage(toolCallInfo *message.EventToolCallInfo, handler message.StreamFunc) {
if gt.active { if mt.active {
gt.endGroup(handler) mt.endMessage(handler)
} }
gt.active = true mt.active = true
gt.groupID = fmt.Sprintf("grp_tool_%d", time.Now().UnixNano()) // Generate message ID using context's ID generator
gt.groupType = message.ChunkToolCall if mt.idGenerator != nil {
gt.startTime = time.Now().UnixMilli() mt.messageID = mt.idGenerator.GenerateMessageID() // M1, M2, M3...
gt.chunkCount = 0 } else {
gt.toolCallInfo = toolCallInfo // Fallback to global generator if no context generator
mt.messageID = message.GenerateNanoID()
}
mt.messageType = message.ChunkToolCall
mt.startTime = time.Now().UnixMilli()
mt.chunkCount = 0
mt.toolCallInfo = toolCallInfo
if handler != nil { if handler != nil {
startData := &message.GroupStartData{ startData := &message.EventMessageStartData{
GroupID: gt.groupID, MessageID: mt.messageID,
Type: string(message.ChunkToolCall), Type: string(message.ChunkToolCall),
Timestamp: gt.startTime, Timestamp: mt.startTime,
ToolCall: toolCallInfo, ToolCall: toolCallInfo,
} }
if startJSON, err := jsoniter.Marshal(startData); err == nil { if startJSON, err := jsoniter.Marshal(startData); err == nil {
@ -70,39 +84,41 @@ func (gt *groupTracker) startToolCallGroup(toolCallInfo *message.GroupToolCallIn
} }
} }
// incrementChunk increments the chunk count for the current group // incrementChunk increments the chunk count for the current message
func (gt *groupTracker) incrementChunk() { func (mt *messageTracker) incrementChunk() {
if gt.active { if mt.active {
gt.chunkCount++ mt.chunkCount++
} }
} }
// endGroup ends the current group and sends group_end event // endMessage ends the current message and sends group_end event
func (gt *groupTracker) endGroup(handler message.StreamFunc) { // Note: group_end event is used for backward compatibility
if !gt.active { // but at LLM level it represents message completion, not Agent-level block
func (mt *messageTracker) endMessage(handler message.StreamFunc) {
if !mt.active {
return return
} }
if handler != nil { if handler != nil {
endData := &message.GroupEndData{ endData := &message.EventMessageEndData{
GroupID: gt.groupID, MessageID: mt.messageID,
Type: string(gt.groupType), Type: string(mt.messageType),
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),
DurationMs: time.Now().UnixMilli() - gt.startTime, DurationMs: time.Now().UnixMilli() - mt.startTime,
ChunkCount: gt.chunkCount, ChunkCount: mt.chunkCount,
Status: "completed", Status: "completed",
} }
if gt.toolCallInfo != nil { if mt.toolCallInfo != nil {
endData.ToolCall = gt.toolCallInfo endData.ToolCall = mt.toolCallInfo
} }
if endJSON, err := jsoniter.Marshal(endData); err == nil { if endJSON, err := jsoniter.Marshal(endData); err == nil {
handler(message.ChunkGroupEnd, endJSON) handler(message.ChunkGroupEnd, endJSON)
} }
} }
gt.active = false mt.active = false
gt.groupID = "" mt.messageID = ""
gt.toolCallInfo = nil mt.toolCallInfo = nil
} }
// Provider OpenAI-compatible provider with capability adapters // Provider OpenAI-compatible provider with capability adapters
@ -438,8 +454,10 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
toolCalls: make(map[int]*accumulatedToolCall), toolCalls: make(map[int]*accumulatedToolCall),
} }
// Group tracker for lifecycle events // Message tracker for lifecycle events (tracks individual messages like thinking, text, tool_call)
groupTracker := &groupTracker{} messageTracker := &messageTracker{
idGenerator: ctx.IDGenerator,
}
// Stream handler // Stream handler
streamHandler := func(data []byte) int { streamHandler := func(data []byte) int {
@ -518,43 +536,43 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
// Handle reasoning content (DeepSeek R1) // Handle reasoning content (DeepSeek R1)
if delta.ReasoningContent != "" { if delta.ReasoningContent != "" {
// Start thinking group if not active // Start thinking message if not active
if !groupTracker.active || groupTracker.groupType != message.ChunkThinking { if !messageTracker.active || messageTracker.messageType != message.ChunkThinking {
groupTracker.startGroup(message.ChunkThinking, handler) messageTracker.startMessage(message.ChunkThinking, handler)
} }
accumulator.reasoningContent += delta.ReasoningContent accumulator.reasoningContent += delta.ReasoningContent
if handler != nil { if handler != nil {
handler(message.ChunkThinking, []byte(delta.ReasoningContent)) handler(message.ChunkThinking, []byte(delta.ReasoningContent))
groupTracker.incrementChunk() messageTracker.incrementChunk()
} }
} }
// Handle content // Handle content
if delta.Content != "" { if delta.Content != "" {
// Start text group if not active // Start text message if not active
if !groupTracker.active || groupTracker.groupType != message.ChunkText { if !messageTracker.active || messageTracker.messageType != message.ChunkText {
groupTracker.startGroup(message.ChunkText, handler) messageTracker.startMessage(message.ChunkText, handler)
} }
accumulator.content += delta.Content accumulator.content += delta.Content
if handler != nil { if handler != nil {
handler(message.ChunkText, []byte(delta.Content)) handler(message.ChunkText, []byte(delta.Content))
groupTracker.incrementChunk() messageTracker.incrementChunk()
} }
} }
// Handle refusal // Handle refusal
if delta.Refusal != "" { if delta.Refusal != "" {
// Start refusal group if not active // Start refusal message if not active
if !groupTracker.active || groupTracker.groupType != message.ChunkRefusal { if !messageTracker.active || messageTracker.messageType != message.ChunkRefusal {
groupTracker.startGroup(message.ChunkRefusal, handler) messageTracker.startMessage(message.ChunkRefusal, handler)
} }
accumulator.refusal += delta.Refusal accumulator.refusal += delta.Refusal
if handler != nil { if handler != nil {
handler(message.ChunkRefusal, []byte(delta.Refusal)) handler(message.ChunkRefusal, []byte(delta.Refusal))
groupTracker.incrementChunk() messageTracker.incrementChunk()
} }
} }
@ -564,14 +582,14 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
if _, exists := accumulator.toolCalls[tc.Index]; !exists { if _, exists := accumulator.toolCalls[tc.Index]; !exists {
accumulator.toolCalls[tc.Index] = &accumulatedToolCall{} accumulator.toolCalls[tc.Index] = &accumulatedToolCall{}
// Start new tool call group when we first see this tool call // Start new tool call message when we first see this tool call
if tc.ID != "" { if tc.ID != "" {
toolCallInfo := &message.GroupToolCallInfo{ toolCallInfo := &message.EventToolCallInfo{
ID: tc.ID, ID: tc.ID,
Name: tc.Function.Name, // May be partial or empty initially Name: tc.Function.Name, // May be partial or empty initially
Index: tc.Index, Index: tc.Index,
} }
groupTracker.startToolCallGroup(toolCallInfo, handler) messageTracker.startToolCallMessage(toolCallInfo, handler)
} }
} }
accTC := accumulator.toolCalls[tc.Index] accTC := accumulator.toolCalls[tc.Index]
@ -585,15 +603,15 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
if tc.Function.Name != "" { if tc.Function.Name != "" {
accTC.functionName = tc.Function.Name accTC.functionName = tc.Function.Name
// Update tool call info in tracker // Update tool call info in tracker
if groupTracker.active && groupTracker.toolCallInfo != nil { if messageTracker.active && messageTracker.toolCallInfo != nil {
groupTracker.toolCallInfo.Name = tc.Function.Name messageTracker.toolCallInfo.Name = tc.Function.Name
} }
} }
if tc.Function.Arguments != "" { if tc.Function.Arguments != "" {
accTC.functionArgs += tc.Function.Arguments accTC.functionArgs += tc.Function.Arguments
// Update tool call info in tracker // Update tool call info in tracker
if groupTracker.active && groupTracker.toolCallInfo != nil { if messageTracker.active && messageTracker.toolCallInfo != nil {
groupTracker.toolCallInfo.Arguments = accTC.functionArgs messageTracker.toolCallInfo.Arguments = accTC.functionArgs
} }
} }
} }
@ -602,7 +620,7 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
if handler != nil { if handler != nil {
toolCallData, _ := jsoniter.Marshal(delta.ToolCalls) toolCallData, _ := jsoniter.Marshal(delta.ToolCalls)
handler(message.ChunkToolCall, toolCallData) handler(message.ChunkToolCall, toolCallData)
groupTracker.incrementChunk() messageTracker.incrementChunk()
} }
} }
@ -713,8 +731,8 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
} }
if err != nil { if err != nil {
// End current group if active // End current message if active
groupTracker.endGroup(handler) messageTracker.endMessage(handler)
// Notify handler of error if provided // Notify handler of error if provided
if handler != nil { if handler != nil {
@ -742,8 +760,8 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
err := fmt.Errorf("no data received from OpenAI API") err := fmt.Errorf("no data received from OpenAI API")
// End current group if active // End current message if active
groupTracker.endGroup(handler) messageTracker.endMessage(handler)
// Notify handler of error if provided // Notify handler of error if provided
if handler != nil { if handler != nil {
@ -786,16 +804,16 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess
// Validate tool call results if schema is provided // Validate tool call results if schema is provided
if err := p.validateToolCallResults(options, toolCalls); err != nil { if err := p.validateToolCallResults(options, toolCalls); err != nil {
// End current group // End current message
groupTracker.endGroup(handler) messageTracker.endMessage(handler)
// Tool call validation failed, need to retry with error feedback // Tool call validation failed, need to retry with error feedback
return nil, fmt.Errorf("tool call validation failed: %w", err) return nil, fmt.Errorf("tool call validation failed: %w", err)
} }
} }
// End final group if still active // End final message if still active
groupTracker.endGroup(handler) messageTracker.endMessage(handler)
return response, nil return response, nil
} }

View file

@ -1251,8 +1251,9 @@ func TestOpenAIProxySupport(t *testing.T) {
t.Log("HTTP proxy support is implemented via http.GetTransport using environment variables") t.Log("HTTP proxy support is implemented via http.GetTransport using environment variables")
} }
// TestOpenAIStreamLifecycleEvents tests that LLM-level lifecycle events (group_start/end) are sent correctly // TestOpenAIStreamLifecycleEvents tests that LLM-level lifecycle events are sent correctly
// Note: stream_start/end are now sent at Agent level, not LLM level // LLM layer sends group_start/end for individual messages (thinking, text, tool_call)
// Note: stream_start/end and Agent-level blocks are handled at Agent level
func TestOpenAIStreamLifecycleEvents(t *testing.T) { func TestOpenAIStreamLifecycleEvents(t *testing.T) {
test.Prepare(t, config.Conf) test.Prepare(t, config.Conf)
defer test.Clean() defer test.Clean()
@ -1284,7 +1285,7 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) {
ctx := newTestContext("test-lifecycle", "openai.gpt-4o") ctx := newTestContext("test-lifecycle", "openai.gpt-4o")
// Track lifecycle events (only group-level events at LLM layer) // Track lifecycle events (group_start/end at LLM layer represent message boundaries)
var events []string var events []string
var groupStartReceived, groupEndReceived bool var groupStartReceived, groupEndReceived bool
@ -1300,11 +1301,11 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) {
case message.ChunkGroupStart: case message.ChunkGroupStart:
groupStartReceived = true groupStartReceived = true
var startData message.GroupStartData var startData message.EventMessageStartData
if err := json.Unmarshal(data, &startData); err == nil { if err := json.Unmarshal(data, &startData); err == nil {
t.Logf("✓ group_start: type=%s, group_id=%s", startData.Type, startData.GroupID) t.Logf("✓ group_start (message start): type=%s, id=%s", startData.Type, startData.MessageID)
if startData.GroupID == "" { if startData.MessageID == "" {
t.Error("group_start missing group_id") t.Error("group_start missing message_id")
} }
} else { } else {
t.Errorf("Failed to parse group_start data: %v", err) t.Errorf("Failed to parse group_start data: %v", err)
@ -1312,9 +1313,9 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) {
case message.ChunkGroupEnd: case message.ChunkGroupEnd:
groupEndReceived = true groupEndReceived = true
var endData message.GroupEndData var endData message.EventMessageEndData
if err := json.Unmarshal(data, &endData); err == nil { if err := json.Unmarshal(data, &endData); err == nil {
t.Logf("✓ group_end: type=%s, chunks=%d, duration=%dms", t.Logf("✓ group_end (message end): type=%s, chunks=%d, duration=%dms",
endData.Type, endData.ChunkCount, endData.DurationMs) endData.Type, endData.ChunkCount, endData.DurationMs)
if endData.ChunkCount <= 0 { if endData.ChunkCount <= 0 {
t.Error("group_end should have chunk_count > 0") t.Error("group_end should have chunk_count > 0")
@ -1339,22 +1340,23 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) {
t.Fatal("Response is nil") t.Fatal("Response is nil")
} }
// Validate that LLM-level lifecycle events were received // Validate that LLM-level message lifecycle events were received
if !groupStartReceived { if !groupStartReceived {
t.Error("group_start event was not received") t.Error("group_start (message start) event was not received")
} }
if !groupEndReceived { if !groupEndReceived {
t.Error("group_end event was not received") t.Error("group_end (message end) event was not received")
} }
// Validate event order: group_start should come before group_end // Validate event order: group_start should come before group_end
if len(events) < 2 { if len(events) < 2 {
t.Errorf("Expected at least 2 events (group_start, group_end), got %d", len(events)) t.Errorf("Expected at least 2 events (message start/end), got %d", len(events))
} }
t.Logf("Total events received: %d", len(events)) t.Logf("Total events received: %d", len(events))
t.Log("LLM lifecycle events test completed successfully") t.Log("LLM message lifecycle events test completed successfully")
t.Log("Note: stream_start/end are now tested at Agent level, not LLM level") t.Log("Note: LLM layer group_start/end represent message boundaries (thinking, text, tool_call)")
t.Log(" Agent-level block boundaries and stream_start/end are handled at Agent level")
} }
// TestOpenAIStreamContextCancellation tests that stream respects context cancellation // TestOpenAIStreamContextCancellation tests that stream respects context cancellation

View file

@ -92,12 +92,13 @@ type accumulatedToolCall struct {
functionArgs string functionArgs string
} }
// groupTracker tracks the current group state for lifecycle events // messageTracker tracks the current message state for lifecycle events
type groupTracker struct { type messageTracker struct {
active bool // Whether a group is currently active active bool // Whether a message is currently active
groupID string // Current group ID messageID string // Current message ID
groupType message.StreamChunkType // Current group type messageType message.StreamChunkType // Current message type (thinking, text, tool_call)
startTime int64 // Group start timestamp startTime int64 // Message start timestamp
chunkCount int // Number of chunks in this group chunkCount int // Number of chunks in this message
toolCallInfo *message.GroupToolCallInfo // Tool call info if group is tool_call type toolCallInfo *message.EventToolCallInfo // Tool call info if message is tool_call type
idGenerator *message.IDGenerator // ID generator from context
} }

View file

@ -447,10 +447,10 @@ msg := output.NewEventMessage("stream_start", "Starting stream...", map[string]i
```go ```go
// Send stream start event (automatically generated by assistant) // Send stream start event (automatically generated by assistant)
// This is typically handled by the framework, not manually sent // This is typically handled by the framework, not manually sent
startData := context.StreamStartData{ startData := message.EventStreamStartData{
RequestID: ctx.RequestID(), RequestID: ctx.RequestID,
Timestamp: time.Now().UnixMilli(), Timestamp: time.Now().UnixMilli(),
TraceID: ctx.TraceID(), TraceID: ctx.Stack.TraceID,
ChatID: ctx.ChatID, ChatID: ctx.ChatID,
} }
output.Send(ctx, output.NewEventMessage("stream_start", "Stream started", startData)) output.Send(ctx, output.NewEventMessage("stream_start", "Stream started", startData))
@ -459,9 +459,13 @@ output.Send(ctx, output.NewEventMessage("stream_start", "Stream started", startD
processData() processData()
// Send stream end event // Send stream end event
output.Send(ctx, output.NewEventMessage("stream_end", "Stream completed", map[string]interface{}{ endData := message.EventStreamEndData{
"duration_ms": 1500, RequestID: ctx.RequestID,
})) Timestamp: time.Now().UnixMilli(),
DurationMs: 1500,
Status: "completed",
}
output.Send(ctx, output.NewEventMessage("stream_end", "Stream completed", endData))
``` ```
**Result:** **Result:**

View file

@ -36,22 +36,20 @@ type Message struct {
Type string `json:"type"` // Message type (e.g., "text", "image", "action") Type string `json:"type"` // Message type (e.g., "text", "image", "action")
Props map[string]interface{} `json:"props,omitempty"` // Type-specific properties Props map[string]interface{} `json:"props,omitempty"` // Type-specific properties
// Streaming control // Streaming control - Hierarchical structure for Agent/LLM/MCP streaming
ID string `json:"id,omitempty"` // Unique message ID (for merging in streaming) ChunkID string `json:"chunk_id,omitempty"` // Unique chunk ID (C1, C2, C3...; for dedup/ordering/debugging)
Delta bool `json:"delta,omitempty"` // Whether this is an incremental update MessageID string `json:"message_id,omitempty"` // Logical message ID (M1, M2, M3...; delta merge target; multiple chunks → one message)
BlockID string `json:"block_id,omitempty"` // Block ID (B1, B2, B3...; Agent-level grouping for UI sections)
ThreadID string `json:"thread_id,omitempty"` // Thread ID (T1, T2, T3...; optional; for concurrent streams)
// Delta update control (for incremental props updates) // Delta control
Delta bool `json:"delta,omitempty"` // Whether this is an incremental update
DeltaPath string `json:"delta_path,omitempty"` // Which field to update (e.g., "content", "items.0.name") DeltaPath string `json:"delta_path,omitempty"` // Which field to update (e.g., "content", "items.0.name")
DeltaAction string `json:"delta_action,omitempty"` // How to update ("append", "replace", "merge", "set") DeltaAction string `json:"delta_action,omitempty"` // How to update ("append", "replace", "merge", "set")
// Type correction (for streaming type inference) // Type correction (for streaming type inference)
TypeChange bool `json:"type_change,omitempty"` // Marks this as a type correction message TypeChange bool `json:"type_change,omitempty"` // Marks this as a type correction message
// Message grouping (for semantically related messages)
GroupID string `json:"group_id,omitempty"` // Parent message group ID
GroupStart bool `json:"group_start,omitempty"` // Marks the start of a group
GroupEnd bool `json:"group_end,omitempty"` // Marks the end of a group
// Metadata // Metadata
Metadata *Metadata `json:"metadata,omitempty"` // Timestamp, sequence, trace ID Metadata *Metadata `json:"metadata,omitempty"` // Timestamp, sequence, trace ID
} }
@ -73,18 +71,38 @@ type Message struct {
#### Streaming Control #### Streaming Control
- **`ID`** (optional): Unique identifier for message tracking Hierarchical structure for fine-grained control over streaming in complex Agent/LLM/MCP scenarios:
- Used to merge multiple delta updates into a single message - **`ChunkID`** (optional): Unique chunk identifier
- Auto-generated if not provided
- Example: `"msg_1234567890_9876543210"` - Auto-generated (C1, C2, C3...)
- For deduplication, ordering, and debugging
- Each raw stream fragment gets a unique ChunkID
- **`MessageID`** (optional): Logical message identifier
- Auto-generated (M1, M2, M3...)
- Delta merge target - multiple chunks with same MessageID are merged
- Represents one complete logical message (e.g., one thinking output, one text response)
- Example: `"M1"`
- **`BlockID`** (optional): Output block identifier
- Auto-generated (B1, B2, B3...)
- Agent-level grouping for UI sections
- One LLM call, one MCP call, or one Agent sub-task
- Used for rendering blocks/sections in the UI
- **`ThreadID`** (optional): Thread identifier
- Auto-generated (T1, T2, T3...)
- For concurrent Agent/LLM/MCP calls
- Distinguishes multiple parallel output streams
- **`Delta`** (optional): Marks this as an incremental update - **`Delta`** (optional): Marks this as an incremental update
- `true`: Append/update to existing message with same MessageID
- `true`: Append/update to existing message with same ID
- `false`: Complete message (default) - `false`: Complete message (default)
- Used for streaming LLM responses - Used for streaming LLM responses
- Message completion is signaled via `group_end` event instead
#### Delta Update Control #### Delta Update Control
@ -109,14 +127,6 @@ For complex, structured messages that need field-level updates:
- Frontend should re-render with new type - Frontend should re-render with new type
- Example: Initially sent as `text`, corrected to `thinking` - Example: Initially sent as `text`, corrected to `thinking`
#### Message Grouping
For grouping semantically related messages (e.g., image + caption):
- **`GroupID`** (optional): Identifier for the message group
- **`GroupStart`** (optional): Marks the beginning of a group
- **`GroupEnd`** (optional): Marks the end of a group
#### Metadata #### Metadata
- **`Metadata`** (optional): Additional message metadata - **`Metadata`** (optional): Additional message metadata
@ -146,7 +156,8 @@ For grouping semantically related messages (e.g., image + caption):
```json ```json
// First chunk // First chunk
{ {
"id": "msg_123", "chunk_id": "C1",
"message_id": "M1",
"type": "text", "type": "text",
"delta": true, "delta": true,
"props": { "props": {
@ -156,7 +167,8 @@ For grouping semantically related messages (e.g., image + caption):
// Second chunk (appends) // Second chunk (appends)
{ {
"id": "msg_123", "chunk_id": "C2",
"message_id": "M1",
"type": "text", "type": "text",
"delta": true, "delta": true,
"props": { "props": {
@ -164,16 +176,30 @@ For grouping semantically related messages (e.g., image + caption):
} }
} }
// Final chunk (marks done) // Third chunk
{ {
"id": "msg_123", "chunk_id": "C3",
"message_id": "M1",
"type": "text", "type": "text",
"delta": true, "delta": true,
"done": true,
"props": { "props": {
"content": "!" "content": "!"
} }
} }
// Completion signaled by message_end event (sent separately)
{
"type": "event",
"props": {
"event": "message_end",
"data": {
"message_id": "M1",
"type": "text",
"chunk_count": 3,
"status": "completed"
}
}
}
``` ```
#### Complex Type with Nested Updates #### Complex Type with Nested Updates
@ -181,7 +207,7 @@ For grouping semantically related messages (e.g., image + caption):
```json ```json
// Initial message // Initial message
{ {
"id": "msg_456", "message_id": "M2",
"type": "table", "type": "table",
"props": { "props": {
"columns": ["Name", "Age"], "columns": ["Name", "Age"],
@ -191,7 +217,8 @@ For grouping semantically related messages (e.g., image + caption):
// Add first row // Add first row
{ {
"id": "msg_456", "chunk_id": "C4",
"message_id": "M2",
"type": "table", "type": "table",
"delta": true, "delta": true,
"delta_path": "rows", "delta_path": "rows",
@ -203,7 +230,8 @@ For grouping semantically related messages (e.g., image + caption):
// Add second row // Add second row
{ {
"id": "msg_456", "chunk_id": "C5",
"message_id": "M2",
"type": "table", "type": "table",
"delta": true, "delta": true,
"delta_path": "rows", "delta_path": "rows",
@ -219,7 +247,8 @@ For grouping semantically related messages (e.g., image + caption):
```json ```json
// Initial guess (text) // Initial guess (text)
{ {
"id": "msg_789", "chunk_id": "C6",
"message_id": "M3",
"type": "text", "type": "text",
"delta": true, "delta": true,
"props": { "props": {
@ -229,7 +258,8 @@ For grouping semantically related messages (e.g., image + caption):
// Correction (actually thinking) // Correction (actually thinking)
{ {
"id": "msg_789", "chunk_id": "C7",
"message_id": "M3",
"type": "thinking", "type": "thinking",
"type_change": true, "type_change": true,
"props": { "props": {
@ -238,38 +268,53 @@ For grouping semantically related messages (e.g., image + caption):
} }
``` ```
#### Message Group #### Block Grouping (Agent-level)
```json ```json
// Group start // Block start event
{ {
"group_id": "grp_001", "type": "event",
"group_start": true
}
// Image in group
{
"type": "image",
"group_id": "grp_001",
"props": { "props": {
"url": "https://example.com/photo.jpg", "event": "block_start",
"alt": "Beautiful sunset" "data": {
"block_id": "B1",
"type": "llm",
"label": "Analyzing image"
}
} }
} }
// Caption in group // Thinking message in block
{ {
"message_id": "M4",
"block_id": "B1",
"type": "thinking",
"props": {
"content": "Let me analyze this image..."
}
}
// Text message in block
{
"message_id": "M5",
"block_id": "B1",
"type": "text", "type": "text",
"group_id": "grp_001",
"props": { "props": {
"content": "Captured at Golden Gate Bridge" "content": "This is a beautiful sunset at Golden Gate Bridge"
} }
} }
// Group end // Block end event
{ {
"group_id": "grp_001", "type": "event",
"group_end": true "props": {
"event": "block_end",
"data": {
"block_id": "B1",
"message_count": 2,
"status": "completed"
}
}
} }
``` ```
@ -363,9 +408,13 @@ output.Send(ctx, err)
### Streaming Messages ### Streaming Messages
```go ```go
// Get ID generator from context
idGen := ctx.IDGenerator
// Send delta (incremental) updates // Send delta (incremental) updates
msg := &message.Message{ msg := &message.Message{
ID: "msg_123", ChunkID: idGen.GenerateChunkID(), // C1
MessageID: idGen.GenerateMessageID(), // M1
Type: message.TypeText, Type: message.TypeText,
Delta: true, // Incremental update Delta: true, // Incremental update
Props: map[string]interface{}{ Props: map[string]interface{}{
@ -374,13 +423,21 @@ msg := &message.Message{
} }
output.Send(ctx, msg) output.Send(ctx, msg)
// Send more delta updates... // Send more delta updates (same MessageID for merging)
msg.Props["content"] = " world" msg2 := &message.Message{
output.Send(ctx, msg) ChunkID: idGen.GenerateChunkID(), // C2
MessageID: msg.MessageID, // M1 (same as before)
Type: message.TypeText,
Delta: true,
Props: map[string]interface{}{
"content": " world",
},
}
output.Send(ctx, msg2)
// Mark completion with group_end event // Mark completion with message_end event
endData := message.GroupEndData{ endData := message.EventMessageEndData{
GroupID: "msg_123", MessageID: msg.MessageID, // M1
Type: "text", Type: "text",
Status: "completed", Status: "completed",
ChunkCount: 2, ChunkCount: 2,
@ -388,7 +445,7 @@ endData := message.GroupEndData{
"content": "Hello world!", // Full content "content": "Hello world!", // Full content
}, },
} }
eventMsg := output.NewEventMessage(message.EventGroupEnd, "Group completed", endData) eventMsg := output.NewEventMessage(message.EventMessageEnd, "Message completed", endData)
output.Send(ctx, eventMsg) output.Send(ctx, eventMsg)
``` ```

View file

@ -196,7 +196,7 @@ msg := output.NewTextMessage("Hello world")
```json ```json
{ {
"id": "M2", "message_id": "M2",
"type": "image", "type": "image",
"props": { "props": {
"url": "https://example.com/avatar.jpg" "url": "https://example.com/avatar.jpg"
@ -227,7 +227,7 @@ msg := output.NewTextMessage("Hello world")
```json ```json
{ {
"id": "M3", "message_id": "M3",
"type": "button", "type": "button",
"props": { "props": {
"text": "Approve", "text": "Approve",

View file

@ -59,7 +59,7 @@ func convertText(msg *message.Message, config *AdapterConfig) ([]interface{}, er
content := getStringProp(msg.Props, "content", "") content := getStringProp(msg.Props, "content", "")
return []interface{}{ return []interface{}{
createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"content": content, "content": content,
}), }),
}, nil }, nil
@ -70,7 +70,7 @@ func convertThinking(msg *message.Message, config *AdapterConfig) ([]interface{}
content := getStringProp(msg.Props, "content", "") content := getStringProp(msg.Props, "content", "")
return []interface{}{ return []interface{}{
createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"reasoning_content": content, "reasoning_content": content,
}), }),
}, nil }, nil
@ -83,7 +83,7 @@ func convertLoading(msg *message.Message, config *AdapterConfig) ([]interface{},
// Convert loading to reasoning_content so it shows in OpenAI clients // Convert loading to reasoning_content so it shows in OpenAI clients
return []interface{}{ return []interface{}{
createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"reasoning_content": message, "reasoning_content": message,
}), }),
}, nil }, nil
@ -109,7 +109,7 @@ func convertToolCall(msg *message.Message, config *AdapterConfig) ([]interface{}
} }
return []interface{}{ return []interface{}{
createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"tool_calls": toolCalls, "tool_calls": toolCalls,
}), }),
}, nil }, nil
@ -150,10 +150,10 @@ func convertStreamStart(msg *message.Message, config *AdapterConfig) ([]interfac
return []interface{}{}, nil return []interface{}{}, nil
} }
// Try to convert to StreamStartData // Try to convert to EventStreamStartData
var startData message.StreamStartData var startData message.EventStreamStartData
switch v := data.(type) { switch v := data.(type) {
case message.StreamStartData: case message.EventStreamStartData:
startData = v startData = v
case map[string]interface{}: case map[string]interface{}:
// If it's a map, try to extract traceID // If it's a map, try to extract traceID
@ -192,7 +192,7 @@ func convertStreamStart(msg *message.Message, config *AdapterConfig) ([]interfac
// Convert to thinking format (reasoning_content) // Convert to thinking format (reasoning_content)
// Reasoning models display this as part of the thinking process // Reasoning models display this as part of the thinking process
content := fmt.Sprintf("🔍 %s - [%s](%s)\n", streamStartText, viewTraceText, traceLink) content := fmt.Sprintf("🔍 %s - [%s](%s)\n", streamStartText, viewTraceText, traceLink)
chunk := createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ chunk := createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"reasoning_content": content, "reasoning_content": content,
}) })
return []interface{}{chunk}, nil return []interface{}{chunk}, nil
@ -200,7 +200,7 @@ func convertStreamStart(msg *message.Message, config *AdapterConfig) ([]interfac
// Convert to regular Markdown text // Convert to regular Markdown text
content := fmt.Sprintf("🚀 %s - [%s](%s)\n", streamStartText, viewTraceText, traceLink) content := fmt.Sprintf("🚀 %s - [%s](%s)\n", streamStartText, viewTraceText, traceLink)
chunk := createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ chunk := createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"content": content, "content": content,
}) })
return []interface{}{chunk}, nil return []interface{}{chunk}, nil
@ -228,7 +228,7 @@ func convertImage(msg *message.Message, config *AdapterConfig) ([]interface{}, e
// Transform URL if transformer is provided // Transform URL if transformer is provided
if config.LinkTransformer != nil { if config.LinkTransformer != nil {
transformedURL, err := config.LinkTransformer(url, msg.Type, msg.ID) transformedURL, err := config.LinkTransformer(url, msg.Type, msg.MessageID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -243,7 +243,7 @@ func convertImage(msg *message.Message, config *AdapterConfig) ([]interface{}, e
text := fmt.Sprintf(template, alt, url) text := fmt.Sprintf(template, alt, url)
return []interface{}{ return []interface{}{
createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"content": text, "content": text,
}), }),
}, nil }, nil
@ -271,7 +271,7 @@ func convertToLink(msg *message.Message, config *AdapterConfig) ([]interface{},
} }
return []interface{}{ return []interface{}{
createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{
"content": text, "content": text,
}), }),
}, nil }, nil
@ -283,7 +283,7 @@ func generateViewLink(msg *message.Message, config *AdapterConfig) (string, erro
if url, ok := msg.Props["url"].(string); ok { if url, ok := msg.Props["url"].(string); ok {
// Transform URL if transformer is provided // Transform URL if transformer is provided
if config.LinkTransformer != nil { if config.LinkTransformer != nil {
return config.LinkTransformer(url, msg.Type, msg.ID) return config.LinkTransformer(url, msg.Type, msg.MessageID)
} }
return url, nil return url, nil
} }
@ -294,11 +294,11 @@ func generateViewLink(msg *message.Message, config *AdapterConfig) (string, erro
baseURL = "" // TODO: Get from environment or context baseURL = "" // TODO: Get from environment or context
} }
viewURL := fmt.Sprintf("%s/agent/view/%s/%s", baseURL, msg.Type, msg.ID) viewURL := fmt.Sprintf("%s/agent/view/%s/%s", baseURL, msg.Type, msg.MessageID)
// Transform URL if transformer is provided // Transform URL if transformer is provided
if config.LinkTransformer != nil { if config.LinkTransformer != nil {
return config.LinkTransformer(viewURL, msg.Type, msg.ID) return config.LinkTransformer(viewURL, msg.Type, msg.MessageID)
} }
return viewURL, nil return viewURL, nil

View file

@ -1,19 +1,11 @@
package output package output
import ( import (
"fmt"
"math/rand"
"time"
"github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/agent/output/message"
) )
// Helper functions for creating built-in message types // Helper functions for creating built-in message types
func init() {
rand.Seed(time.Now().UnixNano())
}
// NewUserInputMessage creates a user input message (for frontend display) // NewUserInputMessage creates a user input message (for frontend display)
// content can be string or []ContentPart for multimodal content // content can be string or []ContentPart for multimodal content
func NewUserInputMessage(content interface{}, role, name string) *message.Message { func NewUserInputMessage(content interface{}, role, name string) *message.Message {
@ -150,10 +142,9 @@ func IsBuiltinType(msgType string) bool {
} }
} }
// GenerateID generates a unique message ID // GenerateID generates a unique message ID using nanoid
// Deprecated: Use message.GenerateMessageID(), message.GenerateChunkID(),
// message.GenerateBlockID(), or message.GenerateThreadID() instead
func GenerateID() string { func GenerateID() string {
// Generate a random ID with timestamp prefix for uniqueness return message.GenerateNanoID()
timestamp := time.Now().UnixNano()
random := rand.Int63()
return fmt.Sprintf("msg_%d_%d", timestamp, random)
} }

View file

@ -0,0 +1,317 @@
# Message Streaming Architecture
This document explains the hierarchical streaming architecture for Agent/LLM/MCP message delivery.
## Overview
The streaming system uses a hierarchical structure to handle complex scenarios including:
- Single LLM calls with multiple message types (thinking, tool calls, text)
- Agent logic with multiple sequential operations (LLM → MCP → LLM)
- Concurrent/parallel calls to multiple LLMs or MCPs
- Real-time delta updates for streaming responses
## Hierarchical Structure
```
Agent Stream (entire conversation)
└─ ThreadID (concurrent stream, optional: T1, T2, T3...)
└─ BlockID (output block/section: B1, B2, B3...)
└─ MessageID (logical message: M1, M2, M3...)
└─ ChunkID (stream fragment: C1, C2, C3...)
```
## Field Definitions
### Message Struct Fields
```go
type Message struct {
// Core fields
Type string `json:"type"`
Props map[string]interface{} `json:"props,omitempty"`
// Streaming control
ChunkID string `json:"chunk_id,omitempty"`
MessageID string `json:"message_id,omitempty"`
BlockID string `json:"block_id,omitempty"`
ThreadID string `json:"thread_id,omitempty"`
// Delta control
Delta bool `json:"delta,omitempty"`
DeltaPath string `json:"delta_path,omitempty"`
DeltaAction string `json:"delta_action,omitempty"`
// ...
}
```
### Field Responsibilities
| Field | Generated By | Purpose | Example Values | Required |
| ----------- | -------------------- | ---------------------------------- | ------------------------------------ | ----------------------------------- |
| `ChunkID` | System (auto) | Deduplication, ordering, debugging | `C1`, `C2`, `C3` | Always |
| `MessageID` | LLM Provider/Handler | Delta merge target | `M1`, `M2`, `M3` or `thinking_msg_1` | Required for delta scenarios |
| `BlockID` | Agent Logic | UI block/section rendering | `B1`, `B2`, `B3` or `llm_response_1` | Required when Agent controls blocks |
| `ThreadID` | Agent Logic | Concurrent stream distinction | `T1`, `T2`, `T3` or `thread_llm1` | Optional (concurrent only) |
### Detailed Field Explanation
#### ChunkID (Stream Fragment Identifier)
- **Purpose**: Uniquely identifies each chunk in the stream
- **Generated**: Automatically by the system (sequential: M1, M2, M3...)
- **Used For**:
- Deduplication (prevent duplicate chunks)
- Ordering (maintain correct sequence)
- Debugging (trace message flow)
- **Scope**: Unique within entire Agent stream
- **Always Present**: Yes
**Example:**
```json
{"chunk_id": "C1", "type": "text", "props": {"content": "Hello"}}
{"chunk_id": "C2", "type": "text", "props": {"content": " World"}}
{"chunk_id": "C3", "type": "thinking", "props": {"content": "..."}}
```
#### MessageID (Logical Message Identifier)
- **Purpose**: Groups multiple chunks into one logical message via delta merging
- **Generated**: By LLM Provider or Stream Handler
- **Used For**:
- Delta merge target (frontend merges all chunks with same MessageID)
- Distinguishing different messages within a group
- **Scope**: Unique within a Group
- **Present When**: Delta streaming is used
**Example:**
```json
// Multiple chunks combine into one "thinking" message
{"chunk_id": "C1", "message_id": "M1", "type": "thinking", "props": {"content": "Let me"}, "delta": true}
{"chunk_id": "C2", "message_id": "M1", "type": "thinking", "props": {"content": " think"}, "delta": true}
{"chunk_id": "C3", "message_id": "M1", "type": "thinking", "props": {"content": "..."}, "delta": true}
// Another independent message
{"chunk_id": "C4", "message_id": "M2", "type": "text", "props": {"content": "Hello"}, "delta": true}
```
#### BlockID (Output Block Identifier)
- **Purpose**: Represents one output block/section (e.g., one LLM call, one MCP call)
- **Generated**: By Agent logic
- **Used For**:
- Frontend UI block/section rendering (visual blocks)
- Distinguishing different operations (LLM vs MCP vs custom logic)
- Organizing related messages together
- **Scope**: Unique within entire Agent stream
- **Present When**: Agent explicitly controls output blocks
**Key Concept**: Block represents a semantic unit of work from Agent's perspective, NOT from LLM's perspective. Each block is rendered as a distinct UI section in the frontend.
**Example:**
```json
// BLOCK 1: LLM Response (contains thinking + tool_call + text)
{"chunk_id": "C1", "block_id": "B1", "message_id": "M1", "type": "thinking", ...}
{"chunk_id": "C2", "block_id": "B1", "message_id": "M2", "type": "tool_call", ...}
{"chunk_id": "C3", "block_id": "B1", "message_id": "M3", "type": "text", ...}
// BLOCK 2: MCP Call
{"chunk_id": "C4", "block_id": "B2", "message_id": "M4", "type": "loading", ...}
{"chunk_id": "C5", "block_id": "B2", "message_id": "M5", "type": "text", ...}
// BLOCK 3: Another LLM Response
{"chunk_id": "C6", "block_id": "B3", "message_id": "M6", "type": "text", ...}
```
#### ThreadID (Concurrent Stream Identifier)
- **Purpose**: Distinguishes concurrent/parallel output streams
- **Generated**: By Agent logic when spawning concurrent operations
- **Used For**:
- Separating outputs from parallel LLM/MCP calls
- Maintaining independent streaming contexts
- **Scope**: Unique within entire Agent stream
- **Present When**: Agent makes concurrent calls (optional)
**Example:**
```json
// Main thread
{"chunk_id": "C1", "thread_id": "T1", "block_id": "B1", "message_id": "M1", "type": "text", ...}
// Parallel MCP calls
{"chunk_id": "C2", "thread_id": "T2", "block_id": "B2", "message_id": "M2", "type": "text", ...}
{"chunk_id": "C3", "thread_id": "T3", "block_id": "B3", "message_id": "M3", "type": "text", ...}
```
## Usage Scenarios
### Scenario 1: Simple Text Message
**No streaming, no grouping**
```json
{
"chunk_id": "C1",
"type": "text",
"props": { "content": "Hello World" }
}
```
**Fields Used:**
- `chunk_id`: C1 (auto-generated)
- No `message_id`, `block_id`, or `thread_id` needed
---
### Scenario 2: LLM Streaming Response (Single Message)
**LLM streams one text message**
```json
{"chunk_id": "C1", "message_id": "M1", "type": "text", "props": {"content": "Hello"}, "delta": true}
{"chunk_id": "C2", "message_id": "M1", "type": "text", "props": {"content": " World"}, "delta": true}
{"chunk_id": "C3", "message_id": "M1", "type": "text", "props": {"content": "!"}, "delta": true}
```
**Fields Used:**
- `chunk_id`: C1, C2, C3 (unique per chunk)
- `message_id`: M1 (same for all, merge target)
- `delta`: true
**Frontend Behavior:**
- Merge all chunks with `message_id: "M1"` into one message
- Display: "Hello World!"
---
### Scenario 3: Agent-Controlled LLM Call (One Block)
**Agent wraps LLM response in an output block**
```typescript
// Agent code starts a block for the LLM response
// System generates block_id: "B1"
// LLM returns thinking + tool_call + text
// Agent ends the block
```
```json
// LLM chunks within block B1
{"chunk_id": "C1", "message_id": "M1", "block_id": "B1", "type": "thinking", "props": {...}, "delta": true}
{"chunk_id": "C2", "message_id": "M1", "block_id": "B1", "type": "thinking", "props": {...}, "delta": true}
{"chunk_id": "C3", "message_id": "M2", "block_id": "B1", "type": "tool_call", "props": {...}}
{"chunk_id": "C4", "message_id": "M3", "block_id": "B1", "type": "text", "props": {...}, "delta": true}
{"chunk_id": "C5", "message_id": "M3", "block_id": "B1", "type": "text", "props": {...}, "delta": true}
```
**Fields Used:**
- `chunk_id`: C1~C5 (unique per chunk)
- `message_id`: M1, M2, M3 (per logical message)
- `block_id`: B1 (all belong to same LLM call)
- `delta`: true (for streaming messages)
**Frontend Behavior:**
- Render one block/section for `block_id: "B1"`
- Within this block, show 3 messages:
- Thinking message (chunks C1+C2 merged into M1)
- Tool call message (chunk C3 = M2)
- Text message (chunks C4+C5 merged into M3)
---
### Scenario 4: Agent Sequential Operations (Multiple Blocks)
**Agent orchestrates: LLM → MCP → LLM**
```typescript
// Agent code orchestrates three sequential operations:
// 1. Block B1: First LLM call
// 2. Block B2: MCP call
// 3. Block B3: Second LLM call
```
```json
// BLOCK 1: First LLM call
{"chunk_id": "C1", "message_id": "M1", "block_id": "B1", "type": "text", ...}
{"chunk_id": "C2", "message_id": "M1", "block_id": "B1", "type": "text", ...}
// BLOCK 2: MCP call
{"chunk_id": "C3", "message_id": "M2", "block_id": "B2", "type": "loading", ...}
{"chunk_id": "C4", "message_id": "M3", "block_id": "B2", "type": "text", ...}
// BLOCK 3: Second LLM call
{"chunk_id": "C5", "message_id": "M4", "block_id": "B3", "type": "text", ...}
{"chunk_id": "C6", "message_id": "M4", "block_id": "B3", "type": "text", ...}
```
**Frontend Behavior:**
- Render 3 distinct blocks/sections:
1. Block 1 (B1): LLM response with text
2. Block 2 (B2): MCP call with loading + result
3. Block 3 (B3): LLM response with text
---
### Scenario 5: Concurrent Operations (Blocks + Threads)
**Agent uses concurrent handler to make parallel calls**
```typescript
// Agent orchestrates parallel operations within one block (B1)
// The concurrent handler automatically assigns thread_id to each operation:
// - MCP call for weather (thread_id: "T1")
// - MCP call for news (thread_id: "T2")
// - LLM call for summary (thread_id: "T3")
//
// Messages from different threads may arrive in any order
```
```json
// Same block, different threads (may arrive in any order)
{"chunk_id": "C1", "message_id": "M1", "block_id": "B1", "thread_id": "T1", "type": "text", "props": {"content": "Weather: Sunny"}}
{"chunk_id": "C2", "message_id": "M2", "block_id": "B1", "thread_id": "T2", "type": "text", "props": {"content": "News: ..."}}
{"chunk_id": "C3", "message_id": "M1", "block_id": "B1", "thread_id": "T1", "type": "text", "props": {"content": ", 25°C"}}
{"chunk_id": "C4", "message_id": "M3", "block_id": "B1", "thread_id": "T3", "type": "text", "props": {"content": "Summary..."}}
```
**Fields Used:**
- `chunk_id`: C1, C2, C3, C4 (unique per chunk, chronological order)
- `message_id`: M1, M2, M3 (per operation/message)
- `block_id`: B1 (all belong to same parallel operation block)
- `thread_id`: T1, T2, T3 (distinguish concurrent operations)
**Frontend Behavior:**
- Render one block for `block_id: "B1"`
- Within this block, separate messages by `thread_id`:
- Thread T1 (Weather): M1 (chunks C1+C3 merged) → "Weather: Sunny, 25°C"
- Thread T2 (News): M2 (chunk C2)
- Thread T3 (Summary): M3 (chunk C4)
- Or interleave by `chunk_id` order (C1, C2, C3, C4) to show real-time arrival
---
## Summary
| Field | Level | Purpose | Example |
| ----------- | ----------- | ------------------ | ---------- |
| `ChunkID` | System | Transport/debug | C1, C2, C3 |
| `MessageID` | LLM/Handler | Delta merging | M1, M2, M3 |
| `BlockID` | Agent | UI blocks/sections | B1, B2, B3 |
| `ThreadID` | Agent | Concurrency | T1, T2, T3 |
**Key Insight**: Each field serves a distinct purpose at a specific layer of the architecture. This hierarchical design supports simple single-message scenarios while enabling complex Agent orchestration with concurrent operations. Blocks provide natural UI boundaries for organizing related messages.

View file

@ -36,20 +36,21 @@ type Message struct {
Type string `json:"type"` // Message type (frontend decides how to render) Type string `json:"type"` // Message type (frontend decides how to render)
Props map[string]interface{} `json:"props,omitempty"` // Message properties (passed to frontend component) Props map[string]interface{} `json:"props,omitempty"` // Message properties (passed to frontend component)
// Streaming control // Streaming control - Hierarchical structure for Agent/LLM/MCP streaming
ID string `json:"id,omitempty"` // Unique chunk/message ID (each chunk has unique ID; use group_id for merging) // See STREAMING.md for detailed explanation of the streaming architecture
Delta bool `json:"delta,omitempty"` // Whether this is an incremental update ChunkID string `json:"chunk_id,omitempty"` // Unique chunk ID (auto-generated: C1, C2, C3...; for dedup/ordering/debugging)
MessageID string `json:"message_id,omitempty"` // Logical message ID (delta merge target; multiple chunks combine into one message)
BlockID string `json:"block_id,omitempty"` // Output block ID (Agent-level control: one LLM call, one MCP call, etc.; for UI rendering blocks/sections)
ThreadID string `json:"thread_id,omitempty"` // Thread ID (optional; for concurrent Agent/LLM/MCP calls to distinguish output streams)
// Delta update control // Delta control
Delta bool `json:"delta,omitempty"` // Whether this is an incremental update
DeltaPath string `json:"delta_path,omitempty"` // Update path (e.g., "content", "data", "items.0.name") DeltaPath string `json:"delta_path,omitempty"` // Update path (e.g., "content", "data", "items.0.name")
DeltaAction string `json:"delta_action,omitempty"` // Update action (append, replace, merge, set) DeltaAction string `json:"delta_action,omitempty"` // Update action (append, replace, merge, set)
// Type correction (for streaming scenarios) // Type correction (for streaming scenarios)
TypeChange bool `json:"type_change,omitempty"` // Marks this as a type correction message TypeChange bool `json:"type_change,omitempty"` // Marks this as a type correction message
// Message group
GroupID string `json:"group_id,omitempty"` // Group ID (all delta chunks of same logical message share this; used for merging)
// Metadata // Metadata
Metadata *Metadata `json:"metadata,omitempty"` // Additional metadata Metadata *Metadata `json:"metadata,omitempty"` // Additional metadata
} }
@ -92,11 +93,27 @@ const (
) )
// Event types for TypeEvent messages // Event types for TypeEvent messages
// Hierarchical structure: Stream > Thread > Block > Message > Chunk
const ( const (
// Stream level events (Agent layer - overall conversation stream)
EventStreamStart = "stream_start" // Stream started event EventStreamStart = "stream_start" // Stream started event
EventStreamEnd = "stream_end" // Stream ended event EventStreamEnd = "stream_end" // Stream ended event
EventGroupStart = "group_start" // Message group started event
EventGroupEnd = "group_end" // Message group ended event // Thread level events (optional - for concurrent scenarios)
EventThreadStart = "thread_start" // Thread started event
EventThreadEnd = "thread_end" // Thread ended event
// Block level events (Agent layer - logical output sections)
EventBlockStart = "block_start" // Block started event
EventBlockEnd = "block_end" // Block ended event
// Message level events (LLM layer - individual logical messages)
EventMessageStart = "message_start" // Message started event
EventMessageEnd = "message_end" // Message ended event
// Backward compatibility aliases (kept for transition period)
EventGroupStart = "group_start" // Alias for EventMessageStart
EventGroupEnd = "group_end" // Alias for EventMessageEnd
) )
// Standard Props structures for built-in types // Standard Props structures for built-in types
@ -284,9 +301,9 @@ type CompletionTokensDetails struct {
// They provide a standardized way to communicate stream boundaries and metadata // They provide a standardized way to communicate stream boundaries and metadata
// to the frontend, enabling better UI/UX (progress indicators, timing, etc.). // to the frontend, enabling better UI/UX (progress indicators, timing, etc.).
// StreamStartData represents the data for stream_start event // EventStreamStartData represents the data for stream_start event
// Sent when a streaming request begins // Sent when a streaming request begins
type StreamStartData struct { type EventStreamStartData struct {
ContextID string `json:"context_id"` // Context ID for the response ContextID string `json:"context_id"` // Context ID for the response
RequestID string `json:"request_id"` // Unique identifier for this request RequestID string `json:"request_id"` // Unique identifier for this request
Timestamp int64 `json:"timestamp"` // Unix timestamp when stream started Timestamp int64 `json:"timestamp"` // Unix timestamp when stream started
@ -296,9 +313,9 @@ type StreamStartData struct {
Metadata map[string]interface{} `json:"metadata,omitempty"` // Metadata to pass to the page for CUI context Metadata map[string]interface{} `json:"metadata,omitempty"` // Metadata to pass to the page for CUI context
} }
// StreamEndData represents the data for stream_end event // EventStreamEndData represents the data for stream_end event
// Sent when a streaming request completes (successfully or with error) // Sent when a streaming request completes (successfully or with error)
type StreamEndData struct { type EventStreamEndData struct {
RequestID string `json:"request_id"` // Corresponding request ID RequestID string `json:"request_id"` // Corresponding request ID
ContextID string `json:"context_id"` // Context ID for the response ContextID string `json:"context_id"` // Context ID for the response
TraceID string `json:"trace_id"` // Trace ID being used (e.g., "trace-123") TraceID string `json:"trace_id"` // Trace ID being used (e.g., "trace-123")
@ -310,34 +327,83 @@ type StreamEndData struct {
Metadata map[string]interface{} `json:"metadata,omitempty"` // Metadata to pass to the page for CUI context Metadata map[string]interface{} `json:"metadata,omitempty"` // Metadata to pass to the page for CUI context
} }
// GroupStartData represents the data for group_start event // EventMessageStartData represents the data for message_start event
// Sent when a logical message group begins (text, tool_call, thinking, etc.) // Sent when a logical message begins (text, tool_call, thinking, etc.)
type GroupStartData struct { // LLM layer: Marks the beginning of a single logical message output
GroupID string `json:"group_id"` // Unique identifier for this group type EventMessageStartData struct {
Type string `json:"type"` // Group type: "text" | "thinking" | "tool_call" | "refusal" MessageID string `json:"message_id"` // Message ID (M1, M2, M3...)
Timestamp int64 `json:"timestamp"` // Unix timestamp when group started Type string `json:"type"` // Message type: "text" | "thinking" | "tool_call" | "refusal"
ToolCall *GroupToolCallInfo `json:"tool_call,omitempty"` // Tool call metadata (if type is "tool_call") Timestamp int64 `json:"timestamp"` // Unix timestamp when message started
ToolCall *EventToolCallInfo `json:"tool_call,omitempty"` // Tool call metadata (if type is "tool_call")
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata (for custom providers or future extensions) Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata (for custom providers or future extensions)
} }
// GroupEndData represents the data for group_end event // EventMessageEndData represents the data for message_end event
// Sent when a logical message group completes // Sent when a logical message completes
type GroupEndData struct { // LLM layer: Signals that all chunks for this message have been sent, client should merge and process
GroupID string `json:"group_id"` // Corresponding group ID type EventMessageEndData struct {
Type string `json:"type"` // Group type (same as in group_start) MessageID string `json:"message_id"` // Message ID (M1, M2, M3...)
Timestamp int64 `json:"timestamp"` // Unix timestamp when group ended Type string `json:"type"` // Message type (same as in message_start)
DurationMs int64 `json:"duration_ms"` // Duration of this group in milliseconds Timestamp int64 `json:"timestamp"` // Unix timestamp when message ended
ChunkCount int `json:"chunk_count"` // Number of data chunks in this group DurationMs int64 `json:"duration_ms"` // Duration of this message in milliseconds
ChunkCount int `json:"chunk_count"` // Number of data chunks in this message
Status string `json:"status"` // "completed" | "partial" | "error" Status string `json:"status"` // "completed" | "partial" | "error"
ToolCall *GroupToolCallInfo `json:"tool_call,omitempty"` // Complete tool call info (if type is "tool_call") ToolCall *EventToolCallInfo `json:"tool_call,omitempty"` // Complete tool call info (if type is "tool_call")
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata (e.g., complete content for direct use)
}
// EventToolCallInfo contains tool call information for message events
// Used in both message_start (partial info) and message_end (complete info)
type EventToolCallInfo struct {
ID string `json:"id"` // Tool call ID (e.g., "call_abc123")
Name string `json:"name"` // Function name (may be partial in message_start)
Arguments string `json:"arguments,omitempty"` // Complete arguments (only in message_end)
Index int `json:"index"` // Index in the tool calls array
}
// EventBlockStartData represents the data for block_start event
// Sent when an output block begins (one LLM call, one MCP call, one Agent sub-task, etc.)
// Agent layer: Groups multiple related messages into a logical section
type EventBlockStartData struct {
BlockID string `json:"block_id"` // Block ID (B1, B2, B3...)
Type string `json:"type"` // Block type: "llm" | "mcp" | "agent" | "tool" | "mixed"
Timestamp int64 `json:"timestamp"` // Unix timestamp when block started
Label string `json:"label,omitempty"` // Human-readable label (e.g., "Searching knowledge base", "Calling weather API")
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata
} }
// GroupToolCallInfo contains tool call information for group events // EventBlockEndData represents the data for block_end event
// Used in both group_start (partial info) and group_end (complete info) // Sent when an output block completes
type GroupToolCallInfo struct { // Agent layer: Signals that this logical section is complete
ID string `json:"id"` // Tool call ID (e.g., "call_abc123") type EventBlockEndData struct {
Name string `json:"name"` // Function name (may be partial in group_start) BlockID string `json:"block_id"` // Block ID (B1, B2, B3...)
Arguments string `json:"arguments,omitempty"` // Complete arguments (only in group_end) Type string `json:"type"` // Block type (same as in block_start)
Index int `json:"index"` // Index in the tool calls array Timestamp int64 `json:"timestamp"` // Unix timestamp when block ended
DurationMs int64 `json:"duration_ms"` // Duration of this block in milliseconds
MessageCount int `json:"message_count"` // Number of messages in this block
Status string `json:"status"` // "completed" | "partial" | "error"
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata
}
// EventThreadStartData represents the data for thread_start event
// Sent when a concurrent thread begins (parallel Agent/LLM/MCP calls)
// Used in concurrent scenarios to distinguish multiple parallel output streams
type EventThreadStartData struct {
ThreadID string `json:"thread_id"` // Thread ID (T1, T2, T3...)
Type string `json:"type"` // Thread type: "agent" | "llm" | "mcp" | "tool"
Timestamp int64 `json:"timestamp"` // Unix timestamp when thread started
Label string `json:"label,omitempty"` // Human-readable label (e.g., "Parallel search 1", "Background task")
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata
}
// EventThreadEndData represents the data for thread_end event
// Sent when a concurrent thread completes
type EventThreadEndData struct {
ThreadID string `json:"thread_id"` // Thread ID (T1, T2, T3...)
Type string `json:"type"` // Thread type (same as in thread_start)
Timestamp int64 `json:"timestamp"` // Unix timestamp when thread ended
DurationMs int64 `json:"duration_ms"` // Duration of this thread in milliseconds
BlockCount int `json:"block_count"` // Number of blocks in this thread
Status string `json:"status"` // "completed" | "partial" | "error"
Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata
} }

View file

@ -0,0 +1,90 @@
package message
import (
"fmt"
"sync/atomic"
gonanoid "github.com/matoous/go-nanoid/v2"
)
// IDGenerator generates unique IDs within a context (e.g., one conversation stream)
// Each Context should have its own IDGenerator to ensure IDs are unique within that context
type IDGenerator struct {
chunkCounter uint64
messageCounter uint64
blockCounter uint64
threadCounter uint64
}
// NewIDGenerator creates a new ID generator for a context
func NewIDGenerator() *IDGenerator {
return &IDGenerator{}
}
// GenerateChunkID generates a unique chunk ID with prefix C
// Format: C1, C2, C3...
func (g *IDGenerator) GenerateChunkID() string {
id := atomic.AddUint64(&g.chunkCounter, 1)
return fmt.Sprintf("C%d", id)
}
// GenerateMessageID generates a unique message ID with prefix M
// Format: M1, M2, M3...
func (g *IDGenerator) GenerateMessageID() string {
id := atomic.AddUint64(&g.messageCounter, 1)
return fmt.Sprintf("M%d", id)
}
// GenerateBlockID generates a unique block ID with prefix B
// Format: B1, B2, B3...
func (g *IDGenerator) GenerateBlockID() string {
id := atomic.AddUint64(&g.blockCounter, 1)
return fmt.Sprintf("B%d", id)
}
// GenerateThreadID generates a unique thread ID with prefix T
// Format: T1, T2, T3...
func (g *IDGenerator) GenerateThreadID() string {
id := atomic.AddUint64(&g.threadCounter, 1)
return fmt.Sprintf("T%d", id)
}
// Reset resets all counters (useful for testing)
func (g *IDGenerator) Reset() {
atomic.StoreUint64(&g.chunkCounter, 0)
atomic.StoreUint64(&g.messageCounter, 0)
atomic.StoreUint64(&g.blockCounter, 0)
atomic.StoreUint64(&g.threadCounter, 0)
}
// GetCounters returns current counter values (for debugging/testing)
func (g *IDGenerator) GetCounters() (chunk, message, block, thread uint64) {
return atomic.LoadUint64(&g.chunkCounter),
atomic.LoadUint64(&g.messageCounter),
atomic.LoadUint64(&g.blockCounter),
atomic.LoadUint64(&g.threadCounter)
}
// GenerateNanoID generates a unique ID using nanoid
// Returns a 21-character URL-safe string
// This is a static function that doesn't depend on the generator's counter
func GenerateNanoID() string {
id, err := gonanoid.New()
if err != nil {
// Fallback to timestamp-based ID if nanoid fails
return fmt.Sprintf("id_%d", atomic.AddUint64(new(uint64), 1))
}
return id
}
// GenerateCustomID generates a custom ID with prefix and nanoid
// Format: prefix_nanoid (e.g., "msg_V1StGXR8_Z5jdHi6B-myT")
// This is a static function that doesn't depend on the generator's counter
func GenerateCustomID(prefix string) string {
id, err := gonanoid.New()
if err != nil {
// Fallback to timestamp-based ID
return fmt.Sprintf("%s_%d", prefix, atomic.AddUint64(new(uint64), 1))
}
return fmt.Sprintf("%s_%s", prefix, id)
}

View file

@ -0,0 +1,188 @@
package message
import (
"sync"
"testing"
)
func TestIDGenerator(t *testing.T) {
gen := NewIDGenerator()
t.Run("GenerateChunkID", func(t *testing.T) {
id1 := gen.GenerateChunkID()
id2 := gen.GenerateChunkID()
id3 := gen.GenerateChunkID()
if id1 != "C1" {
t.Errorf("Expected C1, got %s", id1)
}
if id2 != "C2" {
t.Errorf("Expected C2, got %s", id2)
}
if id3 != "C3" {
t.Errorf("Expected C3, got %s", id3)
}
})
t.Run("GenerateMessageID", func(t *testing.T) {
gen := NewIDGenerator()
id1 := gen.GenerateMessageID()
id2 := gen.GenerateMessageID()
id3 := gen.GenerateMessageID()
if id1 != "M1" {
t.Errorf("Expected M1, got %s", id1)
}
if id2 != "M2" {
t.Errorf("Expected M2, got %s", id2)
}
if id3 != "M3" {
t.Errorf("Expected M3, got %s", id3)
}
})
t.Run("GenerateBlockID", func(t *testing.T) {
gen := NewIDGenerator()
id1 := gen.GenerateBlockID()
id2 := gen.GenerateBlockID()
id3 := gen.GenerateBlockID()
if id1 != "B1" {
t.Errorf("Expected B1, got %s", id1)
}
if id2 != "B2" {
t.Errorf("Expected B2, got %s", id2)
}
if id3 != "B3" {
t.Errorf("Expected B3, got %s", id3)
}
})
t.Run("GenerateThreadID", func(t *testing.T) {
gen := NewIDGenerator()
id1 := gen.GenerateThreadID()
id2 := gen.GenerateThreadID()
id3 := gen.GenerateThreadID()
if id1 != "T1" {
t.Errorf("Expected T1, got %s", id1)
}
if id2 != "T2" {
t.Errorf("Expected T2, got %s", id2)
}
if id3 != "T3" {
t.Errorf("Expected T3, got %s", id3)
}
})
t.Run("Reset", func(t *testing.T) {
gen := NewIDGenerator()
gen.GenerateChunkID()
gen.GenerateMessageID()
gen.GenerateBlockID()
gen.GenerateThreadID()
gen.Reset()
chunk, message, block, thread := gen.GetCounters()
if chunk != 0 || message != 0 || block != 0 || thread != 0 {
t.Errorf("Expected all counters to be 0 after reset, got chunk=%d, message=%d, block=%d, thread=%d",
chunk, message, block, thread)
}
// Verify IDs start from 1 again
if id := gen.GenerateChunkID(); id != "C1" {
t.Errorf("Expected C1 after reset, got %s", id)
}
if id := gen.GenerateMessageID(); id != "M1" {
t.Errorf("Expected M1 after reset, got %s", id)
}
})
t.Run("ConcurrentAccess", func(t *testing.T) {
gen := NewIDGenerator()
var wg sync.WaitGroup
count := 100
// Test concurrent chunk ID generation
wg.Add(count)
for i := 0; i < count; i++ {
go func() {
defer wg.Done()
gen.GenerateChunkID()
}()
}
wg.Wait()
chunk, _, _, _ := gen.GetCounters()
if chunk != uint64(count) {
t.Errorf("Expected chunk counter to be %d, got %d", count, chunk)
}
})
t.Run("MultipleGenerators", func(t *testing.T) {
gen1 := NewIDGenerator()
gen2 := NewIDGenerator()
id1 := gen1.GenerateMessageID()
id2 := gen2.GenerateMessageID()
// Both should start from M1
if id1 != "M1" || id2 != "M1" {
t.Errorf("Expected both generators to start from M1, got %s and %s", id1, id2)
}
// Advance gen1
gen1.GenerateMessageID()
gen1.GenerateMessageID()
// gen2 should still be at M1
id2_next := gen2.GenerateMessageID()
if id2_next != "M2" {
t.Errorf("Expected gen2 to be at M2, got %s", id2_next)
}
// gen1 should be at M3
id1_next := gen1.GenerateMessageID()
if id1_next != "M4" {
t.Errorf("Expected gen1 to be at M4, got %s", id1_next)
}
})
}
func TestGenerateNanoID(t *testing.T) {
id1 := GenerateNanoID()
id2 := GenerateNanoID()
// NanoID should be 21 characters by default
if len(id1) != 21 {
t.Errorf("Expected NanoID length to be 21, got %d", len(id1))
}
// IDs should be unique
if id1 == id2 {
t.Error("Expected unique NanoIDs, got duplicates")
}
t.Logf("Generated NanoIDs: %s, %s", id1, id2)
}
func TestGenerateCustomID(t *testing.T) {
id1 := GenerateCustomID("msg")
id2 := GenerateCustomID("evt")
// Should have prefix
if len(id1) < 4 || id1[:4] != "msg_" {
t.Errorf("Expected ID to start with 'msg_', got %s", id1)
}
if len(id2) < 4 || id2[:4] != "evt_" {
t.Errorf("Expected ID to start with 'evt_', got %s", id2)
}
// IDs should be unique
if id1 == id2 {
t.Error("Expected unique custom IDs, got duplicates")
}
t.Logf("Generated custom IDs: %s, %s", id1, id2)
}