opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor message handling for improved streaming and event management

Browse source 
- Updated message structures to replace 'StreamStartData' and 'StreamEndData' with 'EventStreamStartData' and 'EventStreamEndData' for better clarity and consistency.
- Introduced 'EventMessageStartData' and 'EventMessageEndData' to represent individual message lifecycle events, enhancing the granularity of message tracking.
- Refactored the 'streamState' and 'groupTracker' to utilize the new message structures, improving the organization and handling of streaming events.
- Enhanced the context management by integrating an ID generator for unique message identifiers, facilitating better tracking of message sequences.
- Updated documentation and tests to reflect the new message structures and ensure proper functionality across the system.
This commit is contained in:
Max 2025-11-26 18:30:43 +08:00
parent f0495ab990
commit ef49827faf
18 changed files with 1072 additions and 302 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
agent/assistant/agent.go
View file

@ -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(),

48
agent/assistant/handlers/stream.go
View file

@ -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()

10
agent/context/context.go
View file

@ -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 == "" {

56
agent/context/jsapi.go
View file

@ -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

40
agent/context/jsapi_helpers.go
View file

@ -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 {

20
agent/context/types.go
View file

@ -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

168
agent/llm/providers/openai/openai.go
View file

@ -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
}

32
agent/llm/providers/openai/openai_test.go
View file

@ -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

17
agent/llm/providers/openai/types.go
View file

@ -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
}

16
agent/output/BUILTIN_TYPES.md
View file

@ -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:**

181
agent/output/README.md
View file

@ -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)
```

4
agent/output/adapters/openai/README.md
View file

@ -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",

30
agent/output/adapters/openai/converter.go
View file

@ -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

17
agent/output/builtin.go
View file

@ -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()
}

317
agent/output/message/STREAMING.md Normal file
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.

136
agent/output/message/types.go
View file

@ -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
}

90
agent/output/message/utils.go Normal file
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)
}

188
agent/output/message/utils_test.go Normal file
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)
}