diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index e85f3179..551c189e 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -704,7 +704,7 @@ func (ast *Assistant) sendAgentStreamStart(ctx *context.Context, handler message } // Build the start data - startData := message.StreamStartData{ + startData := message.EventStreamStartData{ ContextID: ctx.ID, ChatID: ctx.ChatID, TraceID: ctx.TraceID(), @@ -731,7 +731,7 @@ func (ast *Assistant) sendAgentStreamEnd(ctx *context.Context, handler message.S return } - endData := &message.StreamEndData{ + endData := &message.EventStreamEndData{ RequestID: ctx.RequestID(), ContextID: ctx.ID, Timestamp: time.Now().UnixMilli(), diff --git a/agent/assistant/handlers/stream.go b/agent/assistant/handlers/stream.go index ce28fe70..a433cd49 100644 --- a/agent/assistant/handlers/stream.go +++ b/agent/assistant/handlers/stream.go @@ -82,7 +82,7 @@ type streamState struct { func (s *streamState) handleStreamStart(data []byte) int { // Send event message to indicate stream has started // 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) if err != nil { 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 func (s *streamState) handleGroupStart(data []byte) int { // 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 { log.Error("Failed to unmarshal group start data: %v", err) return 0 } - // Use the group ID from the start data, or generate one if not provided - groupID := startData.GroupID + // Use the message ID from the start data, or generate one if not provided + groupID := startData.MessageID if groupID == "" { groupID = generateMessageID() - startData.GroupID = groupID + startData.MessageID = groupID } // Initialize group state with the correct group ID @@ -138,13 +138,13 @@ func (s *streamState) handleText(data []byte) int { s.messageSeq++ // Send delta message - // - ID: Sequential message ID (e.g., msg_001, msg_002) for readability - // - GroupID: Same for all chunks of this logical message (frontend merges by group_id) + // - ChunkID: Sequential chunk ID (C1, C2, C3...) for this fragment + // - MessageID: Same for all chunks of this logical message (frontend merges by message_id) msg := &message.Message{ - ID: s.generateSequentialID(), // Sequential ID for this chunk - GroupID: s.currentGroupID, // Group ID for merging (all chunks share this) - Type: message.TypeText, - Delta: true, + ChunkID: s.generateSequentialID(), // Sequential ID for this chunk + MessageID: s.currentGroupID, // Message ID for merging (all chunks share this) + Type: message.TypeText, + Delta: true, Props: map[string]interface{}{ "content": string(data), }, @@ -173,13 +173,13 @@ func (s *streamState) handleThinking(data []byte) int { s.messageSeq++ // Send delta message - // - ID: Sequential message ID (e.g., msg_001, msg_002) for readability - // - GroupID: Same for all chunks of this logical message (frontend merges by group_id) + // - ChunkID: Sequential chunk ID (C1, C2, C3...) for this fragment + // - MessageID: Same for all chunks of this logical message (frontend merges by message_id) msg := &message.Message{ - ID: s.generateSequentialID(), // Sequential ID for this chunk - GroupID: s.currentGroupID, // Group ID for merging (all chunks share this) - Type: message.TypeThinking, - Delta: true, + ChunkID: s.generateSequentialID(), // Sequential ID for this chunk + MessageID: s.currentGroupID, // Message ID for merging (all chunks share this) + Type: message.TypeThinking, + Delta: true, Props: map[string]interface{}{ "content": string(data), }, @@ -197,9 +197,9 @@ func (s *streamState) handleToolCall(data []byte) int { // Tool calls are usually complete JSON objects // Parse and send as tool_call message msg := &message.Message{ - ID: generateMessageID(), - Type: message.TypeToolCall, - Delta: true, + MessageID: generateMessageID(), + Type: message.TypeToolCall, + Delta: true, Props: map[string]interface{}{ // TODO: Parse tool call data "raw": string(data), @@ -241,9 +241,9 @@ func (s *streamState) handleGroupEnd(data []byte) int { msgType = message.TypeText // Fallback to text if type not set } - // Build GroupEndData with complete content - endData := message.GroupEndData{ - GroupID: s.currentGroupID, // Use the group ID, not message ID + // Build EventMessageEndData with complete content + endData := message.EventMessageEndData{ + MessageID: s.currentGroupID, // Use the message ID Type: msgType, Timestamp: time.Now().UnixMilli(), DurationMs: durationMs, @@ -271,7 +271,7 @@ func (s *streamState) handleGroupEnd(data []byte) int { // handleStreamEnd handles stream end event func (s *streamState) handleStreamEnd(data []byte) int { // Parse the stream end data - var endData message.StreamEndData + var endData message.EventStreamEndData if err := jsoniter.Unmarshal(data, &endData); err != nil { log.Error("Failed to parse stream_end data: %v", err) s.ctx.Flush() diff --git a/agent/context/context.go b/agent/context/context.go index d559a119..3dbc1553 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -9,6 +9,7 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/plan" "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/trace" @@ -29,10 +30,11 @@ func New(parent context.Context, authorized *types.AuthorizedInfo, chatID, paylo // Validate the client type ctx := Context{ - Context: parent, - ID: generateContextID(), // Generate unique ID for the context - Space: plan.NewMemorySharedSpace(), - ChatID: chatID, + Context: parent, + ID: generateContextID(), // Generate unique ID for the context + Space: plan.NewMemorySharedSpace(), + ChatID: chatID, + IDGenerator: message.NewIDGenerator(), // Initialize ID generator for this context } if payload == "" { diff --git a/agent/context/jsapi.go b/agent/context/jsapi.go index 59528bff..db621039 100644 --- a/agent/context/jsapi.go +++ b/agent/context/jsapi.go @@ -185,9 +185,13 @@ func (ctx *Context) sendMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { return bridge.JsException(v8ctx, "invalid message: "+err.Error()) } - // Generate unique ID if not provided - if msg.ID == "" { - msg.ID = output.GenerateID() + // Generate unique MessageID if not provided + if msg.MessageID == "" { + if ctx.IDGenerator != nil { + msg.MessageID = ctx.IDGenerator.GenerateMessageID() + } else { + msg.MessageID = output.GenerateID() + } } // Call ctx.Send @@ -222,9 +226,13 @@ func (ctx *Context) sendGroupMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { return bridge.JsException(v8ctx, "invalid group: "+err.Error()) } - // Generate group ID if not provided + // Generate block ID if not provided if group.ID == "" { - group.ID = output.GenerateID() + if ctx.IDGenerator != nil { + group.ID = ctx.IDGenerator.GenerateBlockID() + } else { + group.ID = output.GenerateID() + } } // Send group_start event @@ -232,8 +240,8 @@ func (ctx *Context) sendGroupMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { startEvent := output.NewEventMessage( message.EventGroupStart, "Group started", - message.GroupStartData{ - GroupID: group.ID, + message.EventMessageStartData{ + MessageID: group.ID, Type: "mixed", // Mixed types in group 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()) } - // Generate IDs for messages and set group_id + // Generate MessageIDs for messages and set BlockID for _, msg := range group.Messages { - if msg.ID == "" { - msg.ID = output.GenerateID() + if msg.MessageID == "" { + 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( message.EventGroupEnd, "Group completed", - message.GroupEndData{ - GroupID: group.ID, + message.EventMessageEndData{ + MessageID: group.ID, Type: "mixed", Timestamp: time.Now().UnixMilli(), DurationMs: time.Since(startTime).Milliseconds(), @@ -303,20 +315,24 @@ func (ctx *Context) sendGroupStartMethod(iso *v8go.Isolate) *v8go.FunctionTempla groupType = args[0].String() } - // Get or generate group ID + // Get or generate block ID var groupID string if len(args) > 1 && args[1].IsString() { groupID = args[1].String() } else { - groupID = output.GenerateID() + if ctx.IDGenerator != nil { + groupID = ctx.IDGenerator.GenerateBlockID() + } else { + groupID = output.GenerateID() + } } // Send group_start event startEvent := output.NewEventMessage( message.EventGroupStart, "Group started", - message.GroupStartData{ - GroupID: groupID, + message.EventMessageStartData{ + MessageID: groupID, Type: groupType, Timestamp: time.Now().UnixMilli(), }, @@ -361,8 +377,8 @@ func (ctx *Context) sendGroupEndMethod(iso *v8go.Isolate) *v8go.FunctionTemplate endEvent := output.NewEventMessage( message.EventGroupEnd, "Group completed", - message.GroupEndData{ - GroupID: groupID, + message.EventMessageEndData{ + MessageID: groupID, Type: "mixed", Timestamp: time.Now().UnixMilli(), DurationMs: 0, // Duration not tracked at this level diff --git a/agent/context/jsapi_helpers.go b/agent/context/jsapi_helpers.go index c70708ae..9da25246 100644 --- a/agent/context/jsapi_helpers.go +++ b/agent/context/jsapi_helpers.go @@ -51,10 +51,21 @@ func parseMessage(v8ctx *v8go.Context, jsValue *v8go.Value) (*message.Message, e msg.Props = props } - // Optional fields - if id, ok := msgMap["id"].(string); ok { - msg.ID = id + // Optional fields - Streaming control + if chunkID, ok := msgMap["chunk_id"].(string); ok { + 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 { 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 { msg.TypeChange = typeChange } - if groupID, ok := msgMap["group_id"].(string); ok { - msg.GroupID = groupID - } // Metadata (optional) 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 } - // Optional fields - if id, ok := msgMap["id"].(string); ok { - msg.ID = id + // Optional fields - Streaming control + if chunkID, ok := msgMap["chunk_id"].(string); ok { + 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 { 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 { msg.TypeChange = typeChange } - if groupID, ok := msgMap["group_id"].(string); ok { - msg.GroupID = groupID - } // Metadata (optional) if metadataMap, ok := msgMap["metadata"].(map[string]interface{}); ok { diff --git a/agent/context/types.go b/agent/context/types.go index 7e1f2d4e..d1f4a554 100644 --- a/agent/context/types.go +++ b/agent/context/types.go @@ -7,6 +7,7 @@ import ( "github.com/yaoapp/gou/plan" "github.com/yaoapp/gou/store" "github.com/yaoapp/yao/agent/output" + "github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/openapi/oauth/types" traceTypes "github.com/yaoapp/yao/trace/types" ) @@ -198,15 +199,16 @@ type Context struct { // Context context.Context - ID string `json:"id"` // Context ID for external interrupt identification - Space plan.Space `json:"-"` // Shared data space, it will be used to share data between the request and the call - Cache store.Store `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache" - Stack *Stack `json:"-"` // Stack, current active stack of the request - Stacks map[string]*Stack `json:"-"` // Stacks, all stacks in this request (for trace logging) - Writer Writer `json:"-"` // Writer, it will be used to write response data to the client - 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 - output *output.Output `json:"-"` // Output, it will be used to write response data to the client + ID string `json:"id"` // Context ID for external interrupt identification + Space plan.Space `json:"-"` // Shared data space, it will be used to share data between the request and the call + Cache store.Store `json:"-"` // Cache store, it will be used to store the message cache, default is "__yao.agent.cache" + Stack *Stack `json:"-"` // Stack, current active stack of the request + Stacks map[string]*Stack `json:"-"` // Stacks, all stacks in this request (for trace logging) + Writer Writer `json:"-"` // Writer, it will be used to write response data to the client + 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 + 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) Capabilities *ModelCapabilities `json:"-"` // Model capabilities for the current connector diff --git a/agent/llm/providers/openai/openai.go b/agent/llm/providers/openai/openai.go index 4aaffade..0c01fe11 100644 --- a/agent/llm/providers/openai/openai.go +++ b/agent/llm/providers/openai/openai.go @@ -18,25 +18,33 @@ import ( "github.com/yaoapp/yao/utils/jsonschema" ) -// startGroup starts a new group and sends group_start event -func (gt *groupTracker) startGroup(groupType message.StreamChunkType, handler message.StreamFunc) { - if gt.active { - // End previous group first - gt.endGroup(handler) +// startMessage starts a new message and sends group_start event +// Note: group_start/group_end events are used for backward compatibility +// but at LLM level they represent message boundaries, not Agent-level blocks +func (mt *messageTracker) startMessage(messageType message.StreamChunkType, handler message.StreamFunc) { + if mt.active { + // End previous message first + mt.endMessage(handler) } - gt.active = true - gt.groupID = fmt.Sprintf("grp_%d", time.Now().UnixNano()) - gt.groupType = groupType - gt.startTime = time.Now().UnixMilli() - gt.chunkCount = 0 - gt.toolCallInfo = nil + mt.active = true + // Generate message ID using context's ID generator + if mt.idGenerator != nil { + mt.messageID = mt.idGenerator.GenerateMessageID() // M1, M2, M3... + } else { + // 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 { - startData := &message.GroupStartData{ - GroupID: gt.groupID, - Type: string(groupType), - Timestamp: gt.startTime, + startData := &message.EventMessageStartData{ + MessageID: mt.messageID, + Type: string(messageType), + Timestamp: mt.startTime, } if startJSON, err := jsoniter.Marshal(startData); err == nil { 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 -func (gt *groupTracker) startToolCallGroup(toolCallInfo *message.GroupToolCallInfo, handler message.StreamFunc) { - if gt.active { - gt.endGroup(handler) +// startToolCallMessage starts a new tool call message with tool call info +func (mt *messageTracker) startToolCallMessage(toolCallInfo *message.EventToolCallInfo, handler message.StreamFunc) { + if mt.active { + mt.endMessage(handler) } - gt.active = true - gt.groupID = fmt.Sprintf("grp_tool_%d", time.Now().UnixNano()) - gt.groupType = message.ChunkToolCall - gt.startTime = time.Now().UnixMilli() - gt.chunkCount = 0 - gt.toolCallInfo = toolCallInfo + mt.active = true + // Generate message ID using context's ID generator + if mt.idGenerator != nil { + mt.messageID = mt.idGenerator.GenerateMessageID() // M1, M2, M3... + } else { + // 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 { - startData := &message.GroupStartData{ - GroupID: gt.groupID, + startData := &message.EventMessageStartData{ + MessageID: mt.messageID, Type: string(message.ChunkToolCall), - Timestamp: gt.startTime, + Timestamp: mt.startTime, ToolCall: toolCallInfo, } 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 -func (gt *groupTracker) incrementChunk() { - if gt.active { - gt.chunkCount++ +// incrementChunk increments the chunk count for the current message +func (mt *messageTracker) incrementChunk() { + if mt.active { + mt.chunkCount++ } } -// endGroup ends the current group and sends group_end event -func (gt *groupTracker) endGroup(handler message.StreamFunc) { - if !gt.active { +// endMessage ends the current message and sends group_end event +// Note: group_end event is used for backward compatibility +// but at LLM level it represents message completion, not Agent-level block +func (mt *messageTracker) endMessage(handler message.StreamFunc) { + if !mt.active { return } if handler != nil { - endData := &message.GroupEndData{ - GroupID: gt.groupID, - Type: string(gt.groupType), + endData := &message.EventMessageEndData{ + MessageID: mt.messageID, + Type: string(mt.messageType), Timestamp: time.Now().UnixMilli(), - DurationMs: time.Now().UnixMilli() - gt.startTime, - ChunkCount: gt.chunkCount, + DurationMs: time.Now().UnixMilli() - mt.startTime, + ChunkCount: mt.chunkCount, Status: "completed", } - if gt.toolCallInfo != nil { - endData.ToolCall = gt.toolCallInfo + if mt.toolCallInfo != nil { + endData.ToolCall = mt.toolCallInfo } if endJSON, err := jsoniter.Marshal(endData); err == nil { handler(message.ChunkGroupEnd, endJSON) } } - gt.active = false - gt.groupID = "" - gt.toolCallInfo = nil + mt.active = false + mt.messageID = "" + mt.toolCallInfo = nil } // 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), } - // Group tracker for lifecycle events - groupTracker := &groupTracker{} + // Message tracker for lifecycle events (tracks individual messages like thinking, text, tool_call) + messageTracker := &messageTracker{ + idGenerator: ctx.IDGenerator, + } // Stream handler streamHandler := func(data []byte) int { @@ -518,43 +536,43 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess // Handle reasoning content (DeepSeek R1) if delta.ReasoningContent != "" { - // Start thinking group if not active - if !groupTracker.active || groupTracker.groupType != message.ChunkThinking { - groupTracker.startGroup(message.ChunkThinking, handler) + // Start thinking message if not active + if !messageTracker.active || messageTracker.messageType != message.ChunkThinking { + messageTracker.startMessage(message.ChunkThinking, handler) } accumulator.reasoningContent += delta.ReasoningContent if handler != nil { handler(message.ChunkThinking, []byte(delta.ReasoningContent)) - groupTracker.incrementChunk() + messageTracker.incrementChunk() } } // Handle content if delta.Content != "" { - // Start text group if not active - if !groupTracker.active || groupTracker.groupType != message.ChunkText { - groupTracker.startGroup(message.ChunkText, handler) + // Start text message if not active + if !messageTracker.active || messageTracker.messageType != message.ChunkText { + messageTracker.startMessage(message.ChunkText, handler) } accumulator.content += delta.Content if handler != nil { handler(message.ChunkText, []byte(delta.Content)) - groupTracker.incrementChunk() + messageTracker.incrementChunk() } } // Handle refusal if delta.Refusal != "" { - // Start refusal group if not active - if !groupTracker.active || groupTracker.groupType != message.ChunkRefusal { - groupTracker.startGroup(message.ChunkRefusal, handler) + // Start refusal message if not active + if !messageTracker.active || messageTracker.messageType != message.ChunkRefusal { + messageTracker.startMessage(message.ChunkRefusal, handler) } accumulator.refusal += delta.Refusal if handler != nil { 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 { 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 != "" { - toolCallInfo := &message.GroupToolCallInfo{ + toolCallInfo := &message.EventToolCallInfo{ ID: tc.ID, Name: tc.Function.Name, // May be partial or empty initially Index: tc.Index, } - groupTracker.startToolCallGroup(toolCallInfo, handler) + messageTracker.startToolCallMessage(toolCallInfo, handler) } } accTC := accumulator.toolCalls[tc.Index] @@ -585,15 +603,15 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess if tc.Function.Name != "" { accTC.functionName = tc.Function.Name // Update tool call info in tracker - if groupTracker.active && groupTracker.toolCallInfo != nil { - groupTracker.toolCallInfo.Name = tc.Function.Name + if messageTracker.active && messageTracker.toolCallInfo != nil { + messageTracker.toolCallInfo.Name = tc.Function.Name } } if tc.Function.Arguments != "" { accTC.functionArgs += tc.Function.Arguments // Update tool call info in tracker - if groupTracker.active && groupTracker.toolCallInfo != nil { - groupTracker.toolCallInfo.Arguments = accTC.functionArgs + if messageTracker.active && messageTracker.toolCallInfo != nil { + messageTracker.toolCallInfo.Arguments = accTC.functionArgs } } } @@ -602,7 +620,7 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess if handler != nil { toolCallData, _ := jsoniter.Marshal(delta.ToolCalls) 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 { - // End current group if active - groupTracker.endGroup(handler) + // End current message if active + messageTracker.endMessage(handler) // Notify handler of error if provided 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") - // End current group if active - groupTracker.endGroup(handler) + // End current message if active + messageTracker.endMessage(handler) // Notify handler of error if provided 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 if err := p.validateToolCallResults(options, toolCalls); err != nil { - // End current group - groupTracker.endGroup(handler) + // End current message + messageTracker.endMessage(handler) // Tool call validation failed, need to retry with error feedback return nil, fmt.Errorf("tool call validation failed: %w", err) } } - // End final group if still active - groupTracker.endGroup(handler) + // End final message if still active + messageTracker.endMessage(handler) return response, nil } diff --git a/agent/llm/providers/openai/openai_test.go b/agent/llm/providers/openai/openai_test.go index af0aeaca..dd4a7da4 100644 --- a/agent/llm/providers/openai/openai_test.go +++ b/agent/llm/providers/openai/openai_test.go @@ -1251,8 +1251,9 @@ func TestOpenAIProxySupport(t *testing.T) { 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 -// Note: stream_start/end are now sent at Agent level, not LLM level +// TestOpenAIStreamLifecycleEvents tests that LLM-level lifecycle events are sent correctly +// 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) { test.Prepare(t, config.Conf) defer test.Clean() @@ -1284,7 +1285,7 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { 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 groupStartReceived, groupEndReceived bool @@ -1300,11 +1301,11 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { case message.ChunkGroupStart: groupStartReceived = true - var startData message.GroupStartData + var startData message.EventMessageStartData if err := json.Unmarshal(data, &startData); err == nil { - t.Logf("✓ group_start: type=%s, group_id=%s", startData.Type, startData.GroupID) - if startData.GroupID == "" { - t.Error("group_start missing group_id") + t.Logf("✓ group_start (message start): type=%s, id=%s", startData.Type, startData.MessageID) + if startData.MessageID == "" { + t.Error("group_start missing message_id") } } else { t.Errorf("Failed to parse group_start data: %v", err) @@ -1312,9 +1313,9 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { case message.ChunkGroupEnd: groupEndReceived = true - var endData message.GroupEndData + var endData message.EventMessageEndData 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) if endData.ChunkCount <= 0 { t.Error("group_end should have chunk_count > 0") @@ -1339,22 +1340,23 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) { t.Fatal("Response is nil") } - // Validate that LLM-level lifecycle events were received + // Validate that LLM-level message lifecycle events were received if !groupStartReceived { - t.Error("group_start event was not received") + t.Error("group_start (message start) event was not received") } 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 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.Log("LLM lifecycle events test completed successfully") - t.Log("Note: stream_start/end are now tested at Agent level, not LLM level") + t.Log("LLM message lifecycle events test completed successfully") + 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 diff --git a/agent/llm/providers/openai/types.go b/agent/llm/providers/openai/types.go index 0efa811f..c9157b49 100644 --- a/agent/llm/providers/openai/types.go +++ b/agent/llm/providers/openai/types.go @@ -92,12 +92,13 @@ type accumulatedToolCall struct { functionArgs string } -// groupTracker tracks the current group state for lifecycle events -type groupTracker struct { - active bool // Whether a group is currently active - groupID string // Current group ID - groupType message.StreamChunkType // Current group type - startTime int64 // Group start timestamp - chunkCount int // Number of chunks in this group - toolCallInfo *message.GroupToolCallInfo // Tool call info if group is tool_call type +// messageTracker tracks the current message state for lifecycle events +type messageTracker struct { + active bool // Whether a message is currently active + messageID string // Current message ID + messageType message.StreamChunkType // Current message type (thinking, text, tool_call) + startTime int64 // Message start timestamp + chunkCount int // Number of chunks in this message + toolCallInfo *message.EventToolCallInfo // Tool call info if message is tool_call type + idGenerator *message.IDGenerator // ID generator from context } diff --git a/agent/output/BUILTIN_TYPES.md b/agent/output/BUILTIN_TYPES.md index 96a376ff..7a77649b 100644 --- a/agent/output/BUILTIN_TYPES.md +++ b/agent/output/BUILTIN_TYPES.md @@ -447,10 +447,10 @@ msg := output.NewEventMessage("stream_start", "Starting stream...", map[string]i ```go // Send stream start event (automatically generated by assistant) // This is typically handled by the framework, not manually sent -startData := context.StreamStartData{ - RequestID: ctx.RequestID(), +startData := message.EventStreamStartData{ + RequestID: ctx.RequestID, Timestamp: time.Now().UnixMilli(), - TraceID: ctx.TraceID(), + TraceID: ctx.Stack.TraceID, ChatID: ctx.ChatID, } 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() // Send stream end event -output.Send(ctx, output.NewEventMessage("stream_end", "Stream completed", map[string]interface{}{ - "duration_ms": 1500, -})) +endData := message.EventStreamEndData{ + RequestID: ctx.RequestID, + Timestamp: time.Now().UnixMilli(), + DurationMs: 1500, + Status: "completed", +} +output.Send(ctx, output.NewEventMessage("stream_end", "Stream completed", endData)) ``` **Result:** diff --git a/agent/output/README.md b/agent/output/README.md index 3dc8414a..e8f7e2db 100644 --- a/agent/output/README.md +++ b/agent/output/README.md @@ -36,22 +36,20 @@ type Message struct { Type string `json:"type"` // Message type (e.g., "text", "image", "action") Props map[string]interface{} `json:"props,omitempty"` // Type-specific properties - // Streaming control - ID string `json:"id,omitempty"` // Unique message ID (for merging in streaming) - Delta bool `json:"delta,omitempty"` // Whether this is an incremental update + // Streaming control - Hierarchical structure for Agent/LLM/MCP streaming + ChunkID string `json:"chunk_id,omitempty"` // Unique chunk ID (C1, C2, C3...; for dedup/ordering/debugging) + 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") DeltaAction string `json:"delta_action,omitempty"` // How to update ("append", "replace", "merge", "set") // Type correction (for streaming type inference) 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 `json:"metadata,omitempty"` // Timestamp, sequence, trace ID } @@ -73,18 +71,38 @@ type Message struct { #### 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 - - Auto-generated if not provided - - Example: `"msg_1234567890_9876543210"` +- **`ChunkID`** (optional): Unique chunk identifier + + - 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 - - - `true`: Append/update to existing message with same ID + - `true`: Append/update to existing message with same MessageID - `false`: Complete message (default) - Used for streaming LLM responses - - Message completion is signaled via `group_end` event instead #### Delta Update Control @@ -109,14 +127,6 @@ For complex, structured messages that need field-level updates: - Frontend should re-render with new type - 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`** (optional): Additional message metadata @@ -146,7 +156,8 @@ For grouping semantically related messages (e.g., image + caption): ```json // First chunk { - "id": "msg_123", + "chunk_id": "C1", + "message_id": "M1", "type": "text", "delta": true, "props": { @@ -156,7 +167,8 @@ For grouping semantically related messages (e.g., image + caption): // Second chunk (appends) { - "id": "msg_123", + "chunk_id": "C2", + "message_id": "M1", "type": "text", "delta": true, "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", "delta": true, - "done": true, "props": { "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 @@ -181,7 +207,7 @@ For grouping semantically related messages (e.g., image + caption): ```json // Initial message { - "id": "msg_456", + "message_id": "M2", "type": "table", "props": { "columns": ["Name", "Age"], @@ -191,7 +217,8 @@ For grouping semantically related messages (e.g., image + caption): // Add first row { - "id": "msg_456", + "chunk_id": "C4", + "message_id": "M2", "type": "table", "delta": true, "delta_path": "rows", @@ -203,7 +230,8 @@ For grouping semantically related messages (e.g., image + caption): // Add second row { - "id": "msg_456", + "chunk_id": "C5", + "message_id": "M2", "type": "table", "delta": true, "delta_path": "rows", @@ -219,7 +247,8 @@ For grouping semantically related messages (e.g., image + caption): ```json // Initial guess (text) { - "id": "msg_789", + "chunk_id": "C6", + "message_id": "M3", "type": "text", "delta": true, "props": { @@ -229,7 +258,8 @@ For grouping semantically related messages (e.g., image + caption): // Correction (actually thinking) { - "id": "msg_789", + "chunk_id": "C7", + "message_id": "M3", "type": "thinking", "type_change": true, "props": { @@ -238,38 +268,53 @@ For grouping semantically related messages (e.g., image + caption): } ``` -#### Message Group +#### Block Grouping (Agent-level) ```json -// Group start +// Block start event { - "group_id": "grp_001", - "group_start": true -} - -// Image in group -{ - "type": "image", - "group_id": "grp_001", + "type": "event", "props": { - "url": "https://example.com/photo.jpg", - "alt": "Beautiful sunset" + "event": "block_start", + "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", - "group_id": "grp_001", "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", - "group_end": true + "type": "event", + "props": { + "event": "block_end", + "data": { + "block_id": "B1", + "message_count": 2, + "status": "completed" + } + } } ``` @@ -363,24 +408,36 @@ output.Send(ctx, err) ### Streaming Messages ```go +// Get ID generator from context +idGen := ctx.IDGenerator + // Send delta (incremental) updates msg := &message.Message{ - ID: "msg_123", - Type: message.TypeText, - Delta: true, // Incremental update + ChunkID: idGen.GenerateChunkID(), // C1 + MessageID: idGen.GenerateMessageID(), // M1 + Type: message.TypeText, + Delta: true, // Incremental update Props: map[string]interface{}{ "content": "Hello", }, } output.Send(ctx, msg) -// Send more delta updates... -msg.Props["content"] = " world" -output.Send(ctx, msg) +// Send more delta updates (same MessageID for merging) +msg2 := &message.Message{ + 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 -endData := message.GroupEndData{ - GroupID: "msg_123", +// Mark completion with message_end event +endData := message.EventMessageEndData{ + MessageID: msg.MessageID, // M1 Type: "text", Status: "completed", ChunkCount: 2, @@ -388,7 +445,7 @@ endData := message.GroupEndData{ "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) ``` diff --git a/agent/output/adapters/openai/README.md b/agent/output/adapters/openai/README.md index b8e6a777..6cfbd3ec 100644 --- a/agent/output/adapters/openai/README.md +++ b/agent/output/adapters/openai/README.md @@ -196,7 +196,7 @@ msg := output.NewTextMessage("Hello world") ```json { - "id": "M2", + "message_id": "M2", "type": "image", "props": { "url": "https://example.com/avatar.jpg" @@ -227,7 +227,7 @@ msg := output.NewTextMessage("Hello world") ```json { - "id": "M3", + "message_id": "M3", "type": "button", "props": { "text": "Approve", diff --git a/agent/output/adapters/openai/converter.go b/agent/output/adapters/openai/converter.go index 91914a2e..64f3b997 100644 --- a/agent/output/adapters/openai/converter.go +++ b/agent/output/adapters/openai/converter.go @@ -59,7 +59,7 @@ func convertText(msg *message.Message, config *AdapterConfig) ([]interface{}, er content := getStringProp(msg.Props, "content", "") return []interface{}{ - createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ + createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{ "content": content, }), }, nil @@ -70,7 +70,7 @@ func convertThinking(msg *message.Message, config *AdapterConfig) ([]interface{} content := getStringProp(msg.Props, "content", "") return []interface{}{ - createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ + createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{ "reasoning_content": content, }), }, nil @@ -83,7 +83,7 @@ func convertLoading(msg *message.Message, config *AdapterConfig) ([]interface{}, // Convert loading to reasoning_content so it shows in OpenAI clients return []interface{}{ - createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ + createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{ "reasoning_content": message, }), }, nil @@ -109,7 +109,7 @@ func convertToolCall(msg *message.Message, config *AdapterConfig) ([]interface{} } return []interface{}{ - createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ + createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{ "tool_calls": toolCalls, }), }, nil @@ -150,10 +150,10 @@ func convertStreamStart(msg *message.Message, config *AdapterConfig) ([]interfac return []interface{}{}, nil } - // Try to convert to StreamStartData - var startData message.StreamStartData + // Try to convert to EventStreamStartData + var startData message.EventStreamStartData switch v := data.(type) { - case message.StreamStartData: + case message.EventStreamStartData: startData = v case map[string]interface{}: // 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) // Reasoning models display this as part of the thinking process 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, }) return []interface{}{chunk}, nil @@ -200,7 +200,7 @@ func convertStreamStart(msg *message.Message, config *AdapterConfig) ([]interfac // Convert to regular Markdown text 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, }) return []interface{}{chunk}, nil @@ -228,7 +228,7 @@ func convertImage(msg *message.Message, config *AdapterConfig) ([]interface{}, e // Transform URL if transformer is provided 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 { return nil, err } @@ -243,7 +243,7 @@ func convertImage(msg *message.Message, config *AdapterConfig) ([]interface{}, e text := fmt.Sprintf(template, alt, url) return []interface{}{ - createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ + createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{ "content": text, }), }, nil @@ -271,7 +271,7 @@ func convertToLink(msg *message.Message, config *AdapterConfig) ([]interface{}, } return []interface{}{ - createOpenAIChunk(msg.ID, config.Model, map[string]interface{}{ + createOpenAIChunk(msg.MessageID, config.Model, map[string]interface{}{ "content": text, }), }, nil @@ -283,7 +283,7 @@ func generateViewLink(msg *message.Message, config *AdapterConfig) (string, erro if url, ok := msg.Props["url"].(string); ok { // Transform URL if transformer is provided if config.LinkTransformer != nil { - return config.LinkTransformer(url, msg.Type, msg.ID) + return config.LinkTransformer(url, msg.Type, msg.MessageID) } return url, nil } @@ -294,11 +294,11 @@ func generateViewLink(msg *message.Message, config *AdapterConfig) (string, erro 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 if config.LinkTransformer != nil { - return config.LinkTransformer(viewURL, msg.Type, msg.ID) + return config.LinkTransformer(viewURL, msg.Type, msg.MessageID) } return viewURL, nil diff --git a/agent/output/builtin.go b/agent/output/builtin.go index e58c8c6f..de62a233 100644 --- a/agent/output/builtin.go +++ b/agent/output/builtin.go @@ -1,19 +1,11 @@ package output import ( - "fmt" - "math/rand" - "time" - "github.com/yaoapp/yao/agent/output/message" ) // Helper functions for creating built-in message types -func init() { - rand.Seed(time.Now().UnixNano()) -} - // NewUserInputMessage creates a user input message (for frontend display) // content can be string or []ContentPart for multimodal content 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 { - // Generate a random ID with timestamp prefix for uniqueness - timestamp := time.Now().UnixNano() - random := rand.Int63() - return fmt.Sprintf("msg_%d_%d", timestamp, random) + return message.GenerateNanoID() } diff --git a/agent/output/message/STREAMING.md b/agent/output/message/STREAMING.md new file mode 100644 index 00000000..3e5a0f1b --- /dev/null +++ b/agent/output/message/STREAMING.md @@ -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. diff --git a/agent/output/message/types.go b/agent/output/message/types.go index b80a716f..11cc5964 100644 --- a/agent/output/message/types.go +++ b/agent/output/message/types.go @@ -36,20 +36,21 @@ type Message struct { Type string `json:"type"` // Message type (frontend decides how to render) Props map[string]interface{} `json:"props,omitempty"` // Message properties (passed to frontend component) - // Streaming control - ID string `json:"id,omitempty"` // Unique chunk/message ID (each chunk has unique ID; use group_id for merging) - Delta bool `json:"delta,omitempty"` // Whether this is an incremental update + // Streaming control - Hierarchical structure for Agent/LLM/MCP streaming + // See STREAMING.md for detailed explanation of the streaming architecture + 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") DeltaAction string `json:"delta_action,omitempty"` // Update action (append, replace, merge, set) // Type correction (for streaming scenarios) 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 `json:"metadata,omitempty"` // Additional metadata } @@ -92,11 +93,27 @@ const ( ) // Event types for TypeEvent messages +// Hierarchical structure: Stream > Thread > Block > Message > Chunk const ( + // Stream level events (Agent layer - overall conversation stream) EventStreamStart = "stream_start" // Stream started 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 @@ -284,9 +301,9 @@ type CompletionTokensDetails struct { // They provide a standardized way to communicate stream boundaries and metadata // 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 -type StreamStartData struct { +type EventStreamStartData struct { ContextID string `json:"context_id"` // Context ID for the response RequestID string `json:"request_id"` // Unique identifier for this request 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 } -// 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) -type StreamEndData struct { +type EventStreamEndData struct { RequestID string `json:"request_id"` // Corresponding request ID ContextID string `json:"context_id"` // Context ID for the response 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 } -// GroupStartData represents the data for group_start event -// Sent when a logical message group begins (text, tool_call, thinking, etc.) -type GroupStartData struct { - GroupID string `json:"group_id"` // Unique identifier for this group - Type string `json:"type"` // Group type: "text" | "thinking" | "tool_call" | "refusal" - Timestamp int64 `json:"timestamp"` // Unix timestamp when group started - ToolCall *GroupToolCallInfo `json:"tool_call,omitempty"` // Tool call metadata (if type is "tool_call") +// EventMessageStartData represents the data for message_start event +// Sent when a logical message begins (text, tool_call, thinking, etc.) +// LLM layer: Marks the beginning of a single logical message output +type EventMessageStartData struct { + MessageID string `json:"message_id"` // Message ID (M1, M2, M3...) + Type string `json:"type"` // Message type: "text" | "thinking" | "tool_call" | "refusal" + 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) } -// GroupEndData represents the data for group_end event -// Sent when a logical message group completes -type GroupEndData struct { - GroupID string `json:"group_id"` // Corresponding group ID - Type string `json:"type"` // Group type (same as in group_start) - Timestamp int64 `json:"timestamp"` // Unix timestamp when group ended - DurationMs int64 `json:"duration_ms"` // Duration of this group in milliseconds - ChunkCount int `json:"chunk_count"` // Number of data chunks in this group +// EventMessageEndData represents the data for message_end event +// Sent when a logical message completes +// LLM layer: Signals that all chunks for this message have been sent, client should merge and process +type EventMessageEndData struct { + MessageID string `json:"message_id"` // Message ID (M1, M2, M3...) + Type string `json:"type"` // Message type (same as in message_start) + Timestamp int64 `json:"timestamp"` // Unix timestamp when message ended + 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" - ToolCall *GroupToolCallInfo `json:"tool_call,omitempty"` // Complete tool call info (if type is "tool_call") - Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata + 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) } -// GroupToolCallInfo contains tool call information for group events -// Used in both group_start (partial info) and group_end (complete info) -type GroupToolCallInfo struct { +// 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 group_start) - Arguments string `json:"arguments,omitempty"` // Complete arguments (only in group_end) + 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 +} + +// EventBlockEndData represents the data for block_end event +// Sent when an output block completes +// Agent layer: Signals that this logical section is complete +type EventBlockEndData struct { + BlockID string `json:"block_id"` // Block ID (B1, B2, B3...) + Type string `json:"type"` // Block type (same as in block_start) + 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 +} diff --git a/agent/output/message/utils.go b/agent/output/message/utils.go new file mode 100644 index 00000000..51c57f41 --- /dev/null +++ b/agent/output/message/utils.go @@ -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) +} diff --git a/agent/output/message/utils_test.go b/agent/output/message/utils_test.go new file mode 100644 index 00000000..94b9117c --- /dev/null +++ b/agent/output/message/utils_test.go @@ -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) +}