diff --git a/agent/context/JSAPI.md b/agent/context/JSAPI.md index 75a86970..c0bf0732 100644 --- a/agent/context/JSAPI.md +++ b/agent/context/JSAPI.md @@ -605,6 +605,111 @@ ctx.Send({ - IDs are guaranteed to be unique within the same request/stream - ThreadID is usually auto-managed by Stack, manual generation is for advanced use cases +### Lifecycle Management + +#### `ctx.EndBlock(block_id): void` + +Manually sends a `block_end` event for the specified block. Use this to explicitly mark the end of a block. + +**Parameters:** + +- `block_id`: String - The block ID to end + +**Returns:** + +- `void` + +**Example:** + +```javascript +// Create a block for grouped messages +const block_id = ctx.BlockID(); // "B1" + +// Send messages in the block +ctx.Send("Analyzing data...", block_id); +ctx.Send("Processing results...", block_id); +ctx.Send("Complete!", block_id); + +// Manually end the block +ctx.EndBlock(block_id); +``` + +**Block Lifecycle Events:** + +When you send messages with a `block_id`: + +1. **First message**: Automatically sends `block_start` event +2. **Subsequent messages**: No additional block events +3. **Manual end**: Call `ctx.EndBlock(block_id)` to send `block_end` event + +**block_end Event Format:** + +```json +{ + "type": "event", + "props": { + "event": "block_end", + "message": "Block ended", + "data": { + "block_id": "B1", + "timestamp": 1764483531624, + "duration_ms": 1523, + "message_count": 3, + "status": "completed" + } + } +} +``` + +**Notes:** + +- `block_start` is sent automatically when the first message with a new `block_id` is sent +- `block_end` must be called manually via `ctx.EndBlock()` +- You can track multiple blocks simultaneously (each has independent lifecycle) +- Automatically flushes output after sending the event + +**Use Cases:** + +```javascript +// Use case 1: Progress reporting in a block +function Create(ctx, messages) { + const block_id = ctx.BlockID(); + + ctx.Send("Step 1: Analyzing data...", block_id); + // ... analysis logic ... + + ctx.Send("Step 2: Processing results...", block_id); + // ... processing logic ... + + ctx.Send("Step 3: Complete!", block_id); + + // Mark the block as complete + ctx.EndBlock(block_id); + + return { messages }; +} + +// Use case 2: Multiple parallel blocks +function Create(ctx, messages) { + const llm_block = ctx.BlockID(); // "B1" + const mcp_block = ctx.BlockID(); // "B2" + + // LLM output block + ctx.Send("Thinking...", llm_block); + const response = callLLM(); + ctx.Send(response, llm_block); + ctx.EndBlock(llm_block); + + // MCP tool call block + ctx.Send("Fetching data...", mcp_block); + const data = ctx.MCP.CallTool("tool", "method", {}); + ctx.Send(`Found ${data.length} results`, mcp_block); + ctx.EndBlock(mcp_block); + + return { messages }; +} +``` + ### Resource Cleanup #### `ctx.Release()` diff --git a/agent/context/context.go b/agent/context/context.go index 2098ee3a..b250fdb2 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -384,10 +384,11 @@ func (ctx *Context) recordMessageMetadata(msg *message.Message) { return } - ctx.messageMetadata.set(msg.MessageID, &MessageMetadata{ + ctx.messageMetadata.setMessage(msg.MessageID, &MessageMetadata{ MessageID: msg.MessageID, BlockID: msg.BlockID, ThreadID: msg.ThreadID, + Type: msg.Type, }) } @@ -397,5 +398,5 @@ func (ctx *Context) getMessageMetadata(messageID string) *MessageMetadata { if ctx.messageMetadata == nil { return nil } - return ctx.messageMetadata.get(messageID) + return ctx.messageMetadata.getMessage(messageID) } diff --git a/agent/context/jsapi.go b/agent/context/jsapi.go index 7b3b003d..c8161175 100644 --- a/agent/context/jsapi.go +++ b/agent/context/jsapi.go @@ -61,6 +61,9 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) { jsObject.Set("BlockID", ctx.blockIDMethod(v8ctx.Isolate())) jsObject.Set("ThreadID", ctx.threadIDMethod(v8ctx.Isolate())) + // Lifecycle methods + jsObject.Set("EndBlock", ctx.endBlockMethod(v8ctx.Isolate())) + // Set MCP object jsObject.Set("MCP", ctx.newMCPObject(v8ctx.Isolate())) @@ -577,6 +580,40 @@ func (ctx *Context) threadIDMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { }) } +// endBlockMethod implements ctx.EndBlock(block_id) +// Usage: ctx.EndBlock("B1") +// Sends a block_end event for the specified block +// Returns: undefined +func (ctx *Context) endBlockMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + // Validate arguments + if len(args) < 1 { + return bridge.JsException(v8ctx, "EndBlock requires block_id argument") + } + + if !args[0].IsString() { + return bridge.JsException(v8ctx, "block_id must be a string") + } + + blockID := args[0].String() + + // Call ctx.EndBlock + if err := ctx.EndBlock(blockID); err != nil { + return bridge.JsException(v8ctx, "EndBlock failed: "+err.Error()) + } + + // Automatically flush after ending block + if err := ctx.Flush(); err != nil { + return bridge.JsException(v8ctx, "Flush failed: "+err.Error()) + } + + return v8go.Undefined(iso) + }) +} + // sendGroupMethod implements ctx.SendGroup(group) // Usage: ctx.SendGroup({ id: "group1", messages: [...] }) // Automatically generates IDs, sends group_start/group_end events, and flushes output diff --git a/agent/context/jsapi_output_test.go b/agent/context/jsapi_output_test.go index 78af5e72..1ebe1ab7 100644 --- a/agent/context/jsapi_output_test.go +++ b/agent/context/jsapi_output_test.go @@ -773,3 +773,57 @@ func TestJsValueBlockIDInheritance(t *testing.T) { } assert.Equal(t, true, result["success"], "Delta operations should inherit block_id") } + +// TestJsValueEndBlock tests the EndBlock method on Context +func TestJsValueEndBlock(t *testing.T) { + + test.Prepare(t, config.Conf) + defer test.Clean() + + // Setup mock writer + mockWriter := newMockResponseWriter() + + // Use New() to properly initialize messageMetadata + ctxValue := New(context.Background(), nil, "test-chat-id", "") + cxt := &ctxValue + cxt.AssistantID = "test-assistant-id" + cxt.Accept = AcceptWebCUI + cxt.Locale = "en" + cxt.Writer = mockWriter + + // Test EndBlock method + res, err := v8.Call(v8.CallOptions{}, ` + function test(ctx) { + try { + // Create a block and send messages + const block_id = ctx.BlockID(); // "B1" + + ctx.Send("Message 1", block_id); + ctx.Send("Message 2", block_id); + ctx.Send("Message 3", block_id); + + // End the block manually + ctx.EndBlock(block_id); + + return { success: true, block_id: block_id }; + } catch (error) { + return { success: false, error: error.message }; + } + }`, cxt) + if err != nil { + t.Fatalf("Call failed: %v", err) + } + + result, ok := res.(map[string]interface{}) + if !ok { + t.Fatalf("Expected map result, got %T", res) + } + if !result["success"].(bool) { + t.Logf("Error: %v", result["error"]) + } + assert.Equal(t, true, result["success"], "EndBlock should work correctly") + + // Verify that block_end event was sent + output := mockWriter.buffer.String() + assert.Contains(t, output, "block_end", "Output should contain block_end event") +} diff --git a/agent/context/message_events_test.go b/agent/context/message_events_test.go new file mode 100644 index 00000000..71f7fd75 --- /dev/null +++ b/agent/context/message_events_test.go @@ -0,0 +1,117 @@ +package context_test + +import ( + "bytes" + stdContext "context" + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" +) + +func TestMessageLifecycleEvents(t *testing.T) { + // Create a mock response writer + var buf bytes.Buffer + mockWriter := &mockResponseWriter{ + buffer: &buf, + headers: make(http.Header), + } + + // Create context using New() to ensure proper initialization + ctx := context.New(stdContext.Background(), nil, "test-chat", "") + ctx.Accept = context.AcceptWebCUI + ctx.Writer = mockWriter + ctx.AssistantID = "test-assistant" + ctx.Locale = "en" + + // Send a simple text message + err := ctx.Send(&message.Message{ + Type: message.TypeText, + Props: map[string]interface{}{ + "content": "Hello World", + }, + }) + assert.NoError(t, err) + + // Flush to ensure all messages are written + ctx.Flush() + + // Parse output to find events + output := buf.String() + t.Logf("Output:\n%s", output) + + lines := bytes.Split([]byte(output), []byte("\n")) + + var messages []map[string]interface{} + for _, line := range lines { + if len(line) == 0 { + continue + } + // CUI format: data: {...} + if bytes.HasPrefix(line, []byte("data: ")) { + line = bytes.TrimPrefix(line, []byte("data: ")) + } + + var msg map[string]interface{} + if err := json.Unmarshal(line, &msg); err == nil { + messages = append(messages, msg) + t.Logf("Message: type=%s", msg["type"]) + } + } + + // Check for events + hasMessageStart := false + hasMessageEnd := false + hasTextMessage := false + + for _, msg := range messages { + msgType, _ := msg["type"].(string) + + if msgType == "event" { + if props, ok := msg["props"].(map[string]interface{}); ok { + if eventType, ok := props["event"].(string); ok { + t.Logf("Event type: %s", eventType) + if eventType == "message_start" { + hasMessageStart = true + t.Logf("✓ Found message_start event") + } + if eventType == "message_end" { + hasMessageEnd = true + t.Logf("✓ Found message_end event: %+v", props) + } + } + } + } else if msgType == "text" { + hasTextMessage = true + t.Logf("✓ Found text message") + } + } + + t.Logf("Summary: start=%v, text=%v, end=%v", hasMessageStart, hasTextMessage, hasMessageEnd) + + assert.True(t, hasMessageStart, "Should have message_start event") + assert.True(t, hasTextMessage, "Should have text message") + assert.True(t, hasMessageEnd, "Should have message_end event") +} + +// mockResponseWriter implements http.ResponseWriter for testing +type mockResponseWriter struct { + buffer *bytes.Buffer + statusCode int + headers http.Header +} + +func (m *mockResponseWriter) Header() http.Header { + return m.headers +} + +func (m *mockResponseWriter) Write(data []byte) (int, error) { + return m.buffer.Write(data) +} + +func (m *mockResponseWriter) WriteHeader(statusCode int) { + m.statusCode = statusCode +} diff --git a/agent/context/openapi.go b/agent/context/openapi.go index e9c19176..38b0e81d 100644 --- a/agent/context/openapi.go +++ b/agent/context/openapi.go @@ -9,9 +9,7 @@ import ( "strings" "github.com/gin-gonic/gin" - "github.com/yaoapp/gou/plan" "github.com/yaoapp/gou/store" - "github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/openapi/oauth/authorized" ) @@ -45,30 +43,26 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest clientType := getClientType(userAgent) clientIP := c.ClientIP() - // Create context with unique ID - ctx := &Context{ - Context: c.Request.Context(), - ID: generateContextID(), - Space: plan.NewMemorySharedSpace(), - Cache: cache, - Writer: c.Writer, - Authorized: authInfo, - ChatID: chatID, - AssistantID: assistantID, - Locale: GetLocale(c, completionReq), - Theme: GetTheme(c, completionReq), - Referer: GetReferer(c, completionReq), - Accept: GetAccept(c, completionReq), - Client: Client{ - Type: clientType, - UserAgent: userAgent, - IP: clientIP, - }, - Route: GetRoute(c, completionReq), - Metadata: GetMetadata(c, completionReq), - Skip: GetSkip(c, completionReq), - IDGenerator: message.NewIDGenerator(), // Initialize context-scoped ID generator + // Create context with unique ID using New() to ensure proper initialization + ctxValue := New(c.Request.Context(), authInfo, chatID, "") + ctx := &ctxValue + + // Set additional fields + ctx.Cache = cache + ctx.Writer = c.Writer + ctx.AssistantID = assistantID + ctx.Locale = GetLocale(c, completionReq) + ctx.Theme = GetTheme(c, completionReq) + ctx.Referer = GetReferer(c, completionReq) + ctx.Accept = GetAccept(c, completionReq) + ctx.Client = Client{ + Type: clientType, + UserAgent: userAgent, + IP: clientIP, } + ctx.Route = GetRoute(c, completionReq) + ctx.Metadata = GetMetadata(c, completionReq) + ctx.Skip = GetSkip(c, completionReq) // Initialize interrupt controller ctx.Interrupt = NewInterruptController() diff --git a/agent/context/output.go b/agent/context/output.go index 3c146baf..e4534c4a 100644 --- a/agent/context/output.go +++ b/agent/context/output.go @@ -1,23 +1,29 @@ package context import ( + "time" + "github.com/yaoapp/yao/agent/output" "github.com/yaoapp/yao/agent/output/message" ) // Send sends a message via the output module -// Automatically manages BlockID, ThreadID, and metadata for delta operations -// - For delta operations: inherits BlockID and ThreadID from original message -// - For new messages: auto-generates BlockID if not specified, sets ThreadID from Stack +// Automatically manages BlockID, ThreadID, lifecycle events, and metadata for delta operations +// - For delta operations: inherits BlockID and ThreadID from original message, increments chunk count +// - For new messages: auto-sets ThreadID from Stack, sends message_start event +// - Sends block_start event when a new BlockID is first encountered // - Records metadata for all sent messages to enable delta inheritance func (ctx *Context) Send(msg *message.Message) error { - output, err := ctx.getOutput() + out, err := ctx.getOutput() if err != nil { return err } - // === Delta operations: Auto-inherit BlockID and ThreadID === - if msg.Delta && msg.MessageID != "" { + // Skip lifecycle events for event-type messages (prevent recursion) + isEventMessage := msg.Type == message.TypeEvent + + // === Delta operations: Auto-inherit and update metadata === + if msg.Delta && msg.MessageID != "" && ctx.messageMetadata != nil { if metadata := ctx.getMessageMetadata(msg.MessageID); metadata != nil { // Inherit BlockID if not specified if msg.BlockID == "" { @@ -27,11 +33,14 @@ func (ctx *Context) Send(msg *message.Message) error { if msg.ThreadID == "" { msg.ThreadID = metadata.ThreadID } + + // Increment chunk count for this message + metadata.ChunkCount++ } } - // === Non-delta operations: Auto-set fields === - if !msg.Delta { + // === Non-delta operations: New message logic === + if !msg.Delta && !isEventMessage { // Auto-set ThreadID for non-root Stack (nested agent calls) if msg.ThreadID == "" && ctx.Stack != nil && !ctx.Stack.IsRoot() { msg.ThreadID = ctx.Stack.ID @@ -40,13 +49,183 @@ func (ctx *Context) Send(msg *message.Message) error { // BlockID is NOT auto-generated by default (only manually specified in special cases) // Example: Send a web card after LLM output, group them in the same Block // Developers can specify via ctx.Send(message, blockId) or message.block_id + + // === Send block_start event if this is a new block === + if msg.BlockID != "" && ctx.messageMetadata != nil { + if ctx.messageMetadata.getBlock(msg.BlockID) == nil { + // New block, send block_start event + blockStartData := message.EventBlockStartData{ + BlockID: msg.BlockID, + Type: "mixed", // Default type, can be enhanced later + Timestamp: time.Now().UnixMilli(), + } + blockStartEvent := output.NewEventMessage(message.EventBlockStart, "Block started", blockStartData) + if err := ctx.sendRaw(blockStartEvent); err != nil { + return err + } + + // Record block metadata + ctx.messageMetadata.setBlock(msg.BlockID, &BlockMetadata{ + BlockID: msg.BlockID, + Type: "mixed", + StartTime: time.Now(), + MessageCount: 0, + }) + } + + // Increment message count for this block + ctx.messageMetadata.updateBlock(msg.BlockID, func(block *BlockMetadata) { + block.MessageCount++ + }) + } + + // === Generate MessageID if not provided === + if msg.MessageID == "" { + if ctx.IDGenerator != nil { + msg.MessageID = ctx.IDGenerator.GenerateMessageID() + } else { + msg.MessageID = message.GenerateNanoID() // Use NanoID generator + } + } + + // === Send message_start event === + messageStartData := message.EventMessageStartData{ + MessageID: msg.MessageID, + Type: msg.Type, + Timestamp: time.Now().UnixMilli(), + } + messageStartEvent := output.NewEventMessage(message.EventMessageStart, "Message started", messageStartData) + if err := ctx.sendRaw(messageStartEvent); err != nil { + return err + } + + // === Record message metadata with start time === + if ctx.messageMetadata != nil { + ctx.messageMetadata.setMessage(msg.MessageID, &MessageMetadata{ + MessageID: msg.MessageID, + BlockID: msg.BlockID, + ThreadID: msg.ThreadID, + Type: msg.Type, + StartTime: time.Now(), + ChunkCount: 1, // Initial chunk + }) + } } - // === Record metadata for subsequent delta operations === - ctx.recordMessageMetadata(msg) - // === Actually send the message === - return output.Send(msg) + if err := out.Send(msg); err != nil { + return err + } + + // === Auto-send message_end for non-delta messages (complete messages) === + if !msg.Delta && !isEventMessage && msg.MessageID != "" && ctx.messageMetadata != nil { + metadata := ctx.messageMetadata.getMessage(msg.MessageID) + if metadata != nil { + // Calculate duration + durationMs := time.Since(metadata.StartTime).Milliseconds() + + // Extract content for the extra field + var content interface{} + if msg.Props != nil { + if c, ok := msg.Props["content"]; ok { + content = c + } + } + + // Build message_end event data + endData := message.EventMessageEndData{ + MessageID: msg.MessageID, + Type: msg.Type, + Timestamp: time.Now().UnixMilli(), + DurationMs: durationMs, + ChunkCount: metadata.ChunkCount, + Status: "completed", + } + + // Add content to extra if available + if content != nil { + endData.Extra = map[string]interface{}{ + "content": content, + } + } + + // Send message_end event + messageEndEvent := output.NewEventMessage(message.EventMessageEnd, "Message completed", endData) + ctx.sendRaw(messageEndEvent) + } + } + + return nil +} + +// EndMessage sends a message_end event for a completed message +// Note: For non-delta messages, message_end is automatically sent by Send() +// This method is primarily for delta streaming scenarios: +// - After all delta chunks are sent for a message, call EndMessage() to finalize it +// - For LLM streaming, this is typically called after receiving ChunkMessageEnd +func (ctx *Context) EndMessage(messageID string, content interface{}) error { + if messageID == "" || ctx.messageMetadata == nil { + return nil + } + + metadata := ctx.messageMetadata.getMessage(messageID) + if metadata == nil { + return nil // Message not found, skip + } + + // Calculate duration + durationMs := time.Since(metadata.StartTime).Milliseconds() + + // Build message_end event data + endData := message.EventMessageEndData{ + MessageID: messageID, + Type: metadata.Type, + Timestamp: time.Now().UnixMilli(), + DurationMs: durationMs, + ChunkCount: metadata.ChunkCount, + Status: "completed", + } + + // Add content to extra if provided + if content != nil { + endData.Extra = map[string]interface{}{ + "content": content, + } + } + + // Send message_end event + messageEndEvent := output.NewEventMessage(message.EventMessageEnd, "Message completed", endData) + return ctx.sendRaw(messageEndEvent) +} + +// EndBlock sends a block_end event for a completed block +// This should be called explicitly when all messages in a block are complete +func (ctx *Context) EndBlock(blockID string) error { + if blockID == "" || ctx.messageMetadata == nil { + return nil + } + + blockMetadata := ctx.messageMetadata.getBlock(blockID) + if blockMetadata == nil { + return nil // Block not found, skip + } + + // Calculate duration + durationMs := time.Since(blockMetadata.StartTime).Milliseconds() + + // Build block_end event data + endData := message.EventBlockEndData{ + BlockID: blockID, + Type: blockMetadata.Type, + Timestamp: time.Now().UnixMilli(), + DurationMs: durationMs, + MessageCount: blockMetadata.MessageCount, + Status: "completed", + } + + // Send block_end event + blockEndEvent := output.NewEventMessage(message.EventBlockEnd, "Block completed", endData) + return ctx.sendRaw(blockEndEvent) } // SendGroup sends a group of messages via the output module @@ -77,6 +256,16 @@ func (ctx *Context) CloseOutput() error { return output.Close() } +// sendRaw sends a message directly without triggering lifecycle events +// Used internally to send event messages without recursion +func (ctx *Context) sendRaw(msg *message.Message) error { + out, err := ctx.getOutput() + if err != nil { + return err + } + return out.Send(msg) +} + // getOutput gets the output writer for the context func (ctx *Context) getOutput() (*output.Output, error) { if ctx.output != nil { diff --git a/agent/context/types.go b/agent/context/types.go index 410cc422..3b973931 100644 --- a/agent/context/types.go +++ b/agent/context/types.go @@ -3,6 +3,7 @@ package context import ( "context" "sync" + "time" "github.com/yaoapp/gou/plan" "github.com/yaoapp/gou/store" @@ -197,36 +198,72 @@ type Skip struct { // MessageMetadata stores metadata for sent messages // Used to inherit BlockID and ThreadID in delta operations type MessageMetadata struct { - MessageID string // Message ID - BlockID string // Block ID - ThreadID string // Thread ID + MessageID string // Message ID + BlockID string // Block ID + ThreadID string // Thread ID + Type string // Message type (text, thinking, etc.) + StartTime time.Time // Message start time (for calculating duration) + ChunkCount int // Number of chunks sent for this message } -// messageMetadataStore provides thread-safe storage for message metadata +// BlockMetadata stores metadata for output blocks +type BlockMetadata struct { + BlockID string // Block ID + Type string // Block type (llm, mcp, agent, etc.) + StartTime time.Time // Block start time + MessageCount int // Number of messages in this block +} + +// messageMetadataStore provides thread-safe storage for message and block metadata type messageMetadataStore struct { - data map[string]*MessageMetadata - mu sync.RWMutex + messages map[string]*MessageMetadata // Message metadata by MessageID + blocks map[string]*BlockMetadata // Block metadata by BlockID + mu sync.RWMutex } // newMessageMetadataStore creates a new message metadata store func newMessageMetadataStore() *messageMetadataStore { return &messageMetadataStore{ - data: make(map[string]*MessageMetadata), + messages: make(map[string]*MessageMetadata), + blocks: make(map[string]*BlockMetadata), } } -// set stores metadata for a message (thread-safe) -func (s *messageMetadataStore) set(messageID string, metadata *MessageMetadata) { +// setMessage stores metadata for a message (thread-safe) +func (s *messageMetadataStore) setMessage(messageID string, metadata *MessageMetadata) { s.mu.Lock() defer s.mu.Unlock() - s.data[messageID] = metadata + s.messages[messageID] = metadata } -// get retrieves metadata for a message (thread-safe) -func (s *messageMetadataStore) get(messageID string) *MessageMetadata { +// getMessage retrieves metadata for a message (thread-safe) +func (s *messageMetadataStore) getMessage(messageID string) *MessageMetadata { s.mu.RLock() defer s.mu.RUnlock() - return s.data[messageID] + return s.messages[messageID] +} + +// setBlock stores metadata for a block (thread-safe) +func (s *messageMetadataStore) setBlock(blockID string, metadata *BlockMetadata) { + s.mu.Lock() + defer s.mu.Unlock() + s.blocks[blockID] = metadata +} + +// getBlock retrieves metadata for a block (thread-safe) +func (s *messageMetadataStore) getBlock(blockID string) *BlockMetadata { + s.mu.RLock() + defer s.mu.RUnlock() + return s.blocks[blockID] +} + +// updateBlock updates block metadata (thread-safe) +func (s *messageMetadataStore) updateBlock(blockID string, update func(*BlockMetadata)) { + s.mu.Lock() + defer s.mu.Unlock() + if block, exists := s.blocks[blockID]; exists { + update(block) + } } // Context the context