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

Implement EndBlock method and enhance message lifecycle management

Browse source 
- Added EndBlock method to explicitly mark the end of a message block, sending a block_end event.
- Updated message handling to include lifecycle events for message start and end, improving tracking and management of message durations.
- Enhanced message metadata structure to support chunk counting and message types, facilitating better message operations.
- Improved JSAPI documentation to include new lifecycle management features, clarifying usage for developers.
This commit is contained in:
Max 2025-11-30 15:09:27 +08:00
parent 1af1afaaee
commit 6c295db8d5
8 changed files with 586 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

105
agent/context/JSAPI.md
View file

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

5
agent/context/context.go
View file

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

37
agent/context/jsapi.go
View file

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

54
agent/context/jsapi_output_test.go
View file

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

117
agent/context/message_events_test.go Normal file
View file

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

44
agent/context/openapi.go
View file

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

213
agent/context/output.go
View file

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

63
agent/context/types.go
View file

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