Implement chat buffer management and step tracking in Assistant

- Introduced methods for initializing and managing a chat buffer, allowing for efficient storage of user inputs and assistant messages during chat sessions.
- Added functionality to track execution steps, including beginning and completing steps, with support for capturing space snapshots and handling errors.
- Enhanced the FlushBuffer method to save buffered messages and steps to the database, ensuring data integrity and recovery capabilities.
- Updated the Stream method to integrate buffer management, ensuring proper handling of chat sessions and message storage.
- Added comprehensive tests to validate buffer initialization, user input handling, and step tracking functionalities.

agent/assistant/agent.go

@@ -57,6 +57,38 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 	_, _, done := context.EnterStack(ctx, ast.ID, opts)
 	defer done()
 
+	// ================================================
+	// Initialize Chat Buffer (for root stack only)
+	// Buffer is flushed in defer block at the end
+	// ================================================
+	ast.InitBuffer(ctx)
+
+	// Track final status for buffer flush
+	var finalStatus = context.StepStatusCompleted
+	var finalError error
+
+	// Defer buffer flush - always executes on exit (success, error, interrupt, panic)
+	defer func() {
+		// Handle panic recovery for status tracking
+		if r := recover(); r != nil {
+			finalStatus = context.ResumeStatusFailed
+			if e, ok := r.(error); ok {
+				finalError = e
+			} else {
+				finalError = fmt.Errorf("panic: %v", r)
+			}
+			log.Error("[AGENT] Panic recovered in Stream: %v", r)
+			// Re-panic after flush to preserve original behavior
+			defer panic(r)
+		}
+
+		// Flush buffer to database
+		ast.FlushBuffer(ctx, finalStatus, finalError)
+	}()
+
+	// Buffer user input messages
+	ast.BufferUserInput(ctx, inputMessages)
+
 	// Determine stream handler
 	streamHandler := ast.getStreamHandler(ctx, opts)
 
@@ -64,6 +96,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 	// so that output adapters can use them when converting stream_start event
 	err = ast.initializeCapabilities(ctx, opts)
 	if err != nil {
+		finalStatus = context.ResumeStatusFailed
+		finalError = err
 		ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 		return nil, err
 	}
@@ -76,6 +110,9 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 	// Use async version to not block the main flow
 	ast.InitializeConversationAsync(ctx, opts)
 
+	// Ensure chat session exists
+	ast.EnsureChat(ctx)
+
 	// Initialize agent trace node
 	agentNode := ast.initAgentTraceNode(ctx, inputMessages)
 
@@ -95,15 +132,27 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 	// Request Create hook ( Optional )
 	var createResponse *context.HookCreateResponse
 	if ast.HookScript != nil {
+		// Begin step tracking for hook_create
+		ast.BeginStep(ctx, context.StepTypeHookCreate, map[string]interface{}{
+			"messages": fullMessages,
+		})
+
 		var err error
 		createResponse, opts, err = ast.HookScript.Create(ctx, fullMessages, opts)
 		if err != nil {
+			finalStatus = context.ResumeStatusFailed
+			finalError = err
 			ast.traceAgentFail(agentNode, err)
 			// Send error stream_end for root stack
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 			return nil, err
 		}
 
+		// Complete step
+		ast.CompleteStep(ctx, map[string]interface{}{
+			"response": createResponse,
+		})
+
 		// Log the create response
 		ast.traceCreateHook(agentNode, createResponse)
 	}
@@ -119,6 +168,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 		// Build the LLM request first
 		completionMessages, completionOptions, err = ast.BuildRequest(ctx, inputMessages, createResponse)
 		if err != nil {
+			finalStatus = context.ResumeStatusFailed
+			finalError = err
 			ast.traceAgentFail(agentNode, err)
 			// Send error stream_end for root stack
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@@ -128,19 +179,34 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 		// Build content - convert extended types (file, data) to standard LLM types (text, image_url, input_audio)
 		completionMessages, err = ast.BuildContent(ctx, completionMessages, completionOptions, opts)
 		if err != nil {
+			finalStatus = context.ResumeStatusFailed
+			finalError = err
 			ast.traceAgentFail(agentNode, err)
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 			return nil, err
 		}
 
+		// Begin step tracking for LLM call
+		ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
+			"messages": completionMessages,
+		})
+
 		// Execute the LLM streaming call
 		completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler, opts)
 		if err != nil {
+			finalStatus = context.ResumeStatusFailed
+			finalError = err
 			ast.traceAgentFail(agentNode, err)
 			// Send error stream_end for root stack
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 			return nil, err
 		}
+
+		// Complete LLM step
+		ast.CompleteStep(ctx, map[string]interface{}{
+			"content":    completionResponse.Content,
+			"tool_calls": completionResponse.ToolCalls,
+		})
 	}
 
 	// ================================================
@@ -155,6 +221,12 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 
 		for attempt := 0; attempt < maxToolRetries; attempt++ {
 
+			// Begin step tracking for tool calls
+			ast.BeginStep(ctx, context.StepTypeTool, map[string]interface{}{
+				"tool_calls": currentResponse.ToolCalls,
+				"attempt":    attempt,
+			})
+
 			// Execute all tool calls
 			toolResults, hasErrors := ast.executeToolCalls(ctx, currentResponse.ToolCalls, attempt)
 
@@ -175,8 +247,11 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 				}
 			}
 
-			// If all successful, break out
+			// If all successful, complete step and break out
 			if !hasErrors {
+				ast.CompleteStep(ctx, map[string]interface{}{
+					"results": toolCallResponses,
+				})
 				log.Trace("[AGENT] All tool calls succeeded (attempt %d)", attempt)
 				break
 			}
@@ -193,6 +268,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			// If no retryable errors, don't retry (MCP internal issues)
 			if !hasRetryableErrors {
 				err := fmt.Errorf("tool calls failed with non-retryable errors (MCP internal issues)")
+				finalStatus = context.ResumeStatusFailed
+				finalError = err
 				log.Error("[AGENT] %v", err)
 				ast.traceAgentFail(agentNode, err)
 				ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@@ -202,19 +279,35 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			// If it's the last attempt, return error
 			if attempt == maxToolRetries-1 {
 				err := fmt.Errorf("tool calls failed after %d attempts", maxToolRetries)
+				finalStatus = context.ResumeStatusFailed
+				finalError = err
 				log.Error("[AGENT] %v", err)
 				ast.traceAgentFail(agentNode, err)
 				ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 				return nil, err
 			}
 
+			// Complete current step (with partial results)
+			ast.CompleteStep(ctx, map[string]interface{}{
+				"results":    toolCallResponses,
+				"has_errors": true,
+			})
+
 			// Build retry messages with tool call results (including errors)
 			retryMessages := ast.buildToolRetryMessages(currentMessages, currentResponse, toolResults)
 
+			// Begin LLM retry step
+			ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
+				"messages":      retryMessages,
+				"retry_attempt": attempt + 1,
+			})
+
 			// Retry LLM call (streaming to keep user informed)
 			log.Trace("[AGENT] Retrying LLM for tool call correction (attempt %d/%d)", attempt+1, maxToolRetries-1)
 			currentResponse, err = ast.executeLLMForToolRetry(ctx, retryMessages, completionOptions, agentNode, streamHandler, opts)
 			if err != nil {
+				finalStatus = context.ResumeStatusFailed
+				finalError = err
 				log.Error("[AGENT] LLM retry failed: %v", err)
 				ast.traceAgentFail(agentNode, err)
 				ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@@ -224,12 +317,20 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			// If LLM didn't return tool calls, it might have given up
 			if currentResponse.ToolCalls == nil {
 				err := fmt.Errorf("LLM did not return tool calls in retry attempt %d", attempt+1)
+				finalStatus = context.ResumeStatusFailed
+				finalError = err
 				log.Error("[AGENT] %v", err)
 				ast.traceAgentFail(agentNode, err)
 				ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 				return nil, err
 			}
 
+			// Complete LLM retry step
+			ast.CompleteStep(ctx, map[string]interface{}{
+				"content":    currentResponse.Content,
+				"tool_calls": currentResponse.ToolCalls,
+			})
+
 			// Update messages for next iteration
 			currentMessages = retryMessages
 		}
@@ -245,6 +346,13 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 	var nextResponse *context.NextHookResponse = nil
 
 	if ast.HookScript != nil {
+		// Begin step tracking for hook_next
+		ast.BeginStep(ctx, context.StepTypeHookNext, map[string]interface{}{
+			"messages":   fullMessages,
+			"completion": completionResponse,
+			"tools":      toolCallResponses,
+		})
+
 		var err error
 		nextResponse, opts, err = ast.HookScript.Next(ctx, &context.NextHookPayload{
 			Messages:   fullMessages,
@@ -252,11 +360,18 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			Tools:      toolCallResponses,
 		}, opts)
 		if err != nil {
+			finalStatus = context.ResumeStatusFailed
+			finalError = err
 			ast.traceAgentFail(agentNode, err)
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 			return nil, err
 		}
 
+		// Complete hook_next step
+		ast.CompleteStep(ctx, map[string]interface{}{
+			"response": nextResponse,
+		})
+
 		// Process Next hook response
 		finalResponse, err = ast.processNextResponse(&NextProcessContext{
 			Context:            ctx,
@@ -268,6 +383,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			CreateResponse:     createResponse,
 		})
 		if err != nil {
+			finalStatus = context.ResumeStatusFailed
+			finalError = err
 			ast.traceAgentFail(agentNode, err)
 			ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
 			return nil, err

agent/assistant/chat.go

@@ -4,9 +4,13 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"time"
 
+	"github.com/google/uuid"
+	"github.com/yaoapp/kun/log"
 	agentcontext "github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/i18n"
+	storetypes "github.com/yaoapp/yao/agent/store/types"
 	"github.com/yaoapp/yao/kb"
 	kbapi "github.com/yaoapp/yao/kb/api"
 	"github.com/yaoapp/yao/trace/types"
@@ -210,6 +214,261 @@ func mergeChatMetadata(defaultMetadata map[string]interface{}, ctx *agentcontext
 	return metadata
 }
 
+// =============================================================================
+// Chat Buffer Integration
+// =============================================================================
+
+// InitBuffer initializes the chat buffer for the context
+// Should be called at the start of Stream() for root stack only
+func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) {
+	// Only initialize for root stack
+	if ctx.Stack == nil || !ctx.Stack.IsRoot() {
+		return
+	}
+
+	// Skip if buffer already exists
+	if ctx.Buffer != nil {
+		return
+	}
+
+	// Skip if History is disabled in options
+	if ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History {
+		log.Trace("[CHAT] Buffer skipped: Skip.History is true")
+		return
+	}
+
+	// Generate request ID if not set
+	requestID := ctx.RequestID()
+	if requestID == "" {
+		requestID = uuid.New().String()
+	}
+
+	ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID)
+	log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s", ctx.ChatID, requestID, ast.ID)
+}
+
+// BufferUserInput adds user input messages to the buffer
+// Should be called after InitBuffer
+func (ast *Assistant) BufferUserInput(ctx *agentcontext.Context, inputMessages []agentcontext.Message) {
+	if ctx.Buffer == nil {
+		return
+	}
+
+	// Convert input messages to buffer format
+	for _, msg := range inputMessages {
+		// Extract content from message
+		var content interface{}
+		var name string
+
+		content = msg.Content
+		if msg.Name != nil {
+			name = *msg.Name
+		}
+
+		ctx.Buffer.AddUserInput(content, name)
+	}
+}
+
+// UpdateSpaceSnapshot updates the space snapshot in the buffer
+// Should be called when space data changes
+func (ast *Assistant) UpdateSpaceSnapshot(ctx *agentcontext.Context) {
+	if ctx.Buffer == nil || ctx.Space == nil {
+		return
+	}
+
+	snapshot := ctx.Space.Snapshot()
+	ctx.Buffer.SetSpaceSnapshot(snapshot)
+}
+
+// BeginStep starts tracking an execution step
+// Returns the step for further updates
+func (ast *Assistant) BeginStep(ctx *agentcontext.Context, stepType string, input map[string]interface{}) *agentcontext.BufferedStep {
+	if ctx.Buffer == nil {
+		return nil
+	}
+
+	// Update space snapshot before beginning step
+	ast.UpdateSpaceSnapshot(ctx)
+
+	return ctx.Buffer.BeginStep(stepType, input, ctx.Stack)
+}
+
+// CompleteStep marks the current step as completed
+func (ast *Assistant) CompleteStep(ctx *agentcontext.Context, output map[string]interface{}) {
+	if ctx.Buffer == nil {
+		return
+	}
+	ctx.Buffer.CompleteStep(output)
+}
+
+// FlushBuffer saves all buffered data to the database
+// Should be called in defer block at the end of Stream()
+func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string, err error) {
+	if ctx.Buffer == nil {
+		return
+	}
+
+	// Only flush for root stack
+	if ctx.Stack == nil || !ctx.Stack.IsRoot() {
+		return
+	}
+
+	// Get chat store
+	chatStore := GetChatStore()
+	if chatStore == nil {
+		log.Error("[CHAT] Chat store not available, cannot flush buffer")
+		return
+	}
+
+	// Mark current step as failed/interrupted if needed
+	if finalStatus != agentcontext.StepStatusCompleted && err != nil {
+		ctx.Buffer.FailCurrentStep(finalStatus, err)
+	}
+
+	// 1. Save all messages (user input + assistant responses)
+	messages := ast.convertBufferedMessages(ctx.Buffer.GetMessages())
+	if len(messages) > 0 {
+		if saveErr := chatStore.SaveMessages(ctx.ChatID, messages); saveErr != nil {
+			log.Error("[CHAT] Failed to save messages: %v", saveErr)
+		} else {
+			log.Trace("[CHAT] Saved %d messages for chat=%s", len(messages), ctx.ChatID)
+		}
+	}
+
+	// 2. Update chat last_message_at
+	if len(messages) > 0 {
+		now := time.Now()
+		if updateErr := chatStore.UpdateChat(ctx.ChatID, map[string]interface{}{
+			"last_message_at": now,
+		}); updateErr != nil {
+			log.Trace("[CHAT] Failed to update last_message_at: %v", updateErr)
+		}
+	}
+
+	// 3. Only save resume steps on error/interrupt (not on success)
+	if finalStatus != agentcontext.StepStatusCompleted {
+		steps := ast.convertBufferedSteps(ctx.Buffer.GetStepsForResume(finalStatus))
+		if len(steps) > 0 {
+			if saveErr := chatStore.SaveResume(steps); saveErr != nil {
+				log.Error("[CHAT] Failed to save resume steps: %v", saveErr)
+			} else {
+				log.Trace("[CHAT] Saved %d resume steps for chat=%s (status=%s)", len(steps), ctx.ChatID, finalStatus)
+			}
+		}
+	}
+}
+
+// convertBufferedMessages converts BufferedMessage slice to store Message slice
+func (ast *Assistant) convertBufferedMessages(buffered []*agentcontext.BufferedMessage) []*storetypes.Message {
+	if len(buffered) == 0 {
+		return nil
+	}
+
+	messages := make([]*storetypes.Message, len(buffered))
+	for i, msg := range buffered {
+		messages[i] = &storetypes.Message{
+			MessageID:   msg.MessageID,
+			ChatID:      msg.ChatID,
+			RequestID:   msg.RequestID,
+			Role:        msg.Role,
+			Type:        msg.Type,
+			Props:       msg.Props,
+			BlockID:     msg.BlockID,
+			ThreadID:    msg.ThreadID,
+			AssistantID: msg.AssistantID,
+			Sequence:    msg.Sequence,
+			Metadata:    msg.Metadata,
+			CreatedAt:   msg.CreatedAt,
+			UpdatedAt:   msg.CreatedAt,
+		}
+	}
+	return messages
+}
+
+// convertBufferedSteps converts BufferedStep slice to store Resume slice
+func (ast *Assistant) convertBufferedSteps(buffered []*agentcontext.BufferedStep) []*storetypes.Resume {
+	if len(buffered) == 0 {
+		return nil
+	}
+
+	steps := make([]*storetypes.Resume, len(buffered))
+	for i, step := range buffered {
+		steps[i] = &storetypes.Resume{
+			ResumeID:      step.ResumeID,
+			ChatID:        step.ChatID,
+			RequestID:     step.RequestID,
+			AssistantID:   step.AssistantID,
+			StackID:       step.StackID,
+			StackParentID: step.StackParentID,
+			StackDepth:    step.StackDepth,
+			Type:          step.Type,
+			Status:        step.Status,
+			Input:         step.Input,
+			Output:        step.Output,
+			SpaceSnapshot: step.SpaceSnapshot,
+			Error:         step.Error,
+			Sequence:      step.Sequence,
+			Metadata:      step.Metadata,
+			CreatedAt:     step.CreatedAt,
+			UpdatedAt:     step.CreatedAt,
+		}
+	}
+	return steps
+}
+
+// EnsureChat ensures a chat session exists, creates if not
+func (ast *Assistant) EnsureChat(ctx *agentcontext.Context) error {
+	if ctx.ChatID == "" {
+		return nil // No chat ID, skip
+	}
+
+	chatStore := GetChatStore()
+	if chatStore == nil {
+		return nil // No store, skip
+	}
+
+	// Check if chat exists
+	_, err := chatStore.GetChat(ctx.ChatID)
+	if err == nil {
+		return nil // Chat exists
+	}
+
+	// Create new chat with permission fields
+	chat := &storetypes.Chat{
+		ChatID:      ctx.ChatID,
+		AssistantID: ast.ID,
+		Mode:        "chat",
+		Status:      "active",
+		Share:       "private",
+		Sort:        0,
+		CreatedAt:   time.Now(),
+		UpdatedAt:   time.Now(),
+	}
+
+	// Set permission fields from authorized info
+	if ctx.Authorized != nil {
+		chat.CreatedBy = ctx.Authorized.UserID
+		chat.UpdatedBy = ctx.Authorized.UserID
+		chat.TeamID = ctx.Authorized.TeamID
+		chat.TenantID = ctx.Authorized.TenantID
+	}
+
+	return chatStore.CreateChat(chat)
+}
+
+// GetChatStore returns the chat store instance
+// Returns nil if storage is not configured
+func GetChatStore() storetypes.ChatStore {
+	if storage == nil {
+		return nil
+	}
+	return storage
+}
+
+// =============================================================================
+// Deprecated methods (kept for compatibility)
+// =============================================================================
+
 func (ast *Assistant) saveChat(ctx *agentcontext.Context, input []agentcontext.Message, opts *agentcontext.Options) error {
 	_ = ctx
 	_ = input

agent/assistant/chat_test.go

@@ -7,10 +7,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/yaoapp/gou/plan"
 	"github.com/yaoapp/yao/agent/assistant"
 	agentcontext "github.com/yaoapp/yao/agent/context"
+	storetypes "github.com/yaoapp/yao/agent/store/types"
 	"github.com/yaoapp/yao/agent/testutils"
 	"github.com/yaoapp/yao/kb"
 	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
@@ -300,3 +303,566 @@ func TestInitializeConversation(t *testing.T) {
 		t.Logf("✓ Correctly skipped with history flag")
 	})
 }
+
+// =============================================================================
+// Buffer Integration Tests
+// =============================================================================
+
+func TestBufferInitialization(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("mohe")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	t.Run("InitBufferForRootStack", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_buffer_001")
+
+		// Enter stack to simulate root stack
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+
+		// Initialize buffer
+		ast.InitBuffer(ctx)
+
+		// Verify buffer was created
+		assert.NotNil(t, ctx.Buffer, "Buffer should be initialized for root stack")
+		assert.Equal(t, "test_chat_buffer_001", ctx.Buffer.ChatID())
+		assert.Equal(t, ast.ID, ctx.Buffer.AssistantID())
+		t.Logf("✓ Buffer initialized: chatID=%s, assistantID=%s", ctx.Buffer.ChatID(), ctx.Buffer.AssistantID())
+	})
+
+	t.Run("SkipBufferForNestedStack", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_buffer_nested")
+
+		// Enter root stack
+		_, _, doneRoot := agentcontext.EnterStack(ctx, "root_assistant", nil)
+		defer doneRoot()
+
+		// Enter nested stack
+		_, _, doneNested := agentcontext.EnterStack(ctx, "nested_assistant", nil)
+		defer doneNested()
+
+		// Try to initialize buffer (should be skipped for nested stack)
+		ast.InitBuffer(ctx)
+
+		// Buffer should be nil because we're not at root
+		assert.Nil(t, ctx.Buffer, "Buffer should not be initialized for nested stack")
+		t.Logf("✓ Buffer correctly skipped for nested stack")
+	})
+
+	t.Run("IdempotentBufferInit", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_buffer_idem")
+
+		// Enter stack
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+
+		// Initialize buffer twice
+		ast.InitBuffer(ctx)
+		firstBuffer := ctx.Buffer
+
+		ast.InitBuffer(ctx)
+		secondBuffer := ctx.Buffer
+
+		// Should be the same buffer instance
+		assert.Same(t, firstBuffer, secondBuffer, "Buffer should be idempotent")
+		t.Logf("✓ Buffer initialization is idempotent")
+	})
+}
+
+func TestBufferUserInput(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("mohe")
+	require.NoError(t, err)
+
+	t.Run("BufferSimpleTextInput", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_input_001")
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Create input messages
+		inputMessages := []agentcontext.Message{
+			{
+				Role:    agentcontext.RoleUser,
+				Content: "Hello, how are you?",
+			},
+		}
+
+		// Buffer user input
+		ast.BufferUserInput(ctx, inputMessages)
+
+		// Verify buffer contains the message
+		messages := ctx.Buffer.GetMessages()
+		assert.Len(t, messages, 1, "Should have 1 buffered message")
+		assert.Equal(t, "user", messages[0].Role)
+		assert.Equal(t, "user_input", messages[0].Type)
+		assert.Equal(t, "Hello, how are you?", messages[0].Props["content"])
+		t.Logf("✓ User input buffered: %v", messages[0].Props)
+	})
+
+	t.Run("BufferMultipleMessages", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_input_multi")
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Create multiple input messages
+		inputMessages := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "First message"},
+			{Role: agentcontext.RoleUser, Content: "Second message"},
+		}
+
+		// Buffer user input
+		ast.BufferUserInput(ctx, inputMessages)
+
+		// Verify buffer contains all messages
+		messages := ctx.Buffer.GetMessages()
+		assert.Len(t, messages, 2, "Should have 2 buffered messages")
+		assert.Equal(t, 1, messages[0].Sequence)
+		assert.Equal(t, 2, messages[1].Sequence)
+		t.Logf("✓ Multiple messages buffered with correct sequence")
+	})
+
+	t.Run("BufferWithNilBuffer", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_input_nil")
+
+		// Don't initialize buffer
+		inputMessages := []agentcontext.Message{
+			{Role: agentcontext.RoleUser, Content: "Test"},
+		}
+
+		// Should not panic
+		ast.BufferUserInput(ctx, inputMessages)
+		t.Logf("✓ BufferUserInput handles nil buffer gracefully")
+	})
+}
+
+func TestBufferStepTracking(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("mohe")
+	require.NoError(t, err)
+
+	t.Run("BeginAndCompleteStep", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_step_001")
+		ctx.Space = plan.NewMemorySharedSpace()
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Set some space data
+		ctx.Space.Set("test_key", "test_value")
+
+		// Begin a step
+		step := ast.BeginStep(ctx, agentcontext.StepTypeLLM, map[string]interface{}{
+			"messages": []string{"Hello"},
+		})
+
+		assert.NotNil(t, step, "Step should be created")
+		assert.Equal(t, agentcontext.StepTypeLLM, step.Type)
+		assert.Equal(t, agentcontext.StepStatusRunning, step.Status)
+		assert.NotEmpty(t, step.StackID)
+
+		// Complete the step
+		ast.CompleteStep(ctx, map[string]interface{}{
+			"content": "Response",
+		})
+
+		// Verify step is completed
+		steps := ctx.Buffer.GetAllSteps()
+		assert.Len(t, steps, 1)
+		assert.Equal(t, agentcontext.StepStatusCompleted, steps[0].Status)
+		assert.Equal(t, "Response", steps[0].Output["content"])
+		t.Logf("✓ Step tracking works correctly")
+	})
+
+	t.Run("SpaceSnapshotCapture", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_space_001")
+		ctx.Space = plan.NewMemorySharedSpace()
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Set space data before step
+		ctx.Space.Set("key1", "value1")
+		ctx.Space.Set("key2", 123)
+
+		// Begin step (should capture space snapshot)
+		ast.BeginStep(ctx, agentcontext.StepTypeHookCreate, nil)
+
+		// Verify space snapshot was captured
+		steps := ctx.Buffer.GetAllSteps()
+		require.Len(t, steps, 1)
+		assert.NotNil(t, steps[0].SpaceSnapshot)
+		assert.Equal(t, "value1", steps[0].SpaceSnapshot["key1"])
+		assert.Equal(t, 123, steps[0].SpaceSnapshot["key2"])
+		t.Logf("✓ Space snapshot captured: %v", steps[0].SpaceSnapshot)
+	})
+
+	t.Run("MultipleSteps", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "test_chat_multi_step")
+		ctx.Space = plan.NewMemorySharedSpace()
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Step 1: hook_create
+		ast.BeginStep(ctx, agentcontext.StepTypeHookCreate, map[string]interface{}{"phase": "create"})
+		ast.CompleteStep(ctx, map[string]interface{}{"result": "created"})
+
+		// Step 2: llm
+		ast.BeginStep(ctx, agentcontext.StepTypeLLM, map[string]interface{}{"phase": "llm"})
+		ast.CompleteStep(ctx, map[string]interface{}{"result": "completed"})
+
+		// Step 3: hook_next
+		ast.BeginStep(ctx, agentcontext.StepTypeHookNext, map[string]interface{}{"phase": "next"})
+		ast.CompleteStep(ctx, map[string]interface{}{"result": "done"})
+
+		// Verify all steps
+		steps := ctx.Buffer.GetAllSteps()
+		assert.Len(t, steps, 3)
+		assert.Equal(t, agentcontext.StepTypeHookCreate, steps[0].Type)
+		assert.Equal(t, agentcontext.StepTypeLLM, steps[1].Type)
+		assert.Equal(t, agentcontext.StepTypeHookNext, steps[2].Type)
+		t.Logf("✓ Multiple steps tracked correctly")
+	})
+}
+
+func TestFlushBuffer(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("mohe")
+	require.NoError(t, err)
+
+	// Skip if chat store not available
+	chatStore := assistant.GetChatStore()
+	if chatStore == nil {
+		t.Skip("Chat store not configured, skipping flush tests")
+	}
+
+	t.Run("FlushOnSuccess", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_flush_success_%s", uuid.New().String()[:8])
+		ctx := agentcontext.New(context.Background(), nil, chatID)
+		ctx.Space = plan.NewMemorySharedSpace()
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Ensure chat exists
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Mode:        "chat",
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+
+		// Add some messages to buffer
+		ctx.Buffer.AddUserInput("Test question", "")
+		ctx.Buffer.AddAssistantMessage("text", map[string]interface{}{"content": "Test answer"}, "", "", ast.ID, nil)
+
+		// Add a step
+		ast.BeginStep(ctx, agentcontext.StepTypeLLM, nil)
+		ast.CompleteStep(ctx, nil)
+
+		// Flush buffer (success case)
+		ast.FlushBuffer(ctx, agentcontext.StepStatusCompleted, nil)
+
+		// Verify messages were saved
+		messages, err := chatStore.GetMessages(chatID, storetypes.MessageFilter{})
+		assert.NoError(t, err)
+		assert.Len(t, messages, 2, "Should have 2 messages saved")
+
+		// Verify no resume records (success case)
+		resumes, err := chatStore.GetResume(chatID)
+		assert.NoError(t, err)
+		assert.Len(t, resumes, 0, "Should have no resume records on success")
+
+		// Cleanup
+		chatStore.DeleteChat(chatID)
+		t.Logf("✓ Buffer flushed on success: %d messages saved, no resume records", len(messages))
+	})
+
+	t.Run("FlushOnFailure", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_flush_fail_%s", uuid.New().String()[:8])
+		ctx := agentcontext.New(context.Background(), nil, chatID)
+		ctx.Space = plan.NewMemorySharedSpace()
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Ensure chat exists
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Mode:        "chat",
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+
+		// Add messages
+		ctx.Buffer.AddUserInput("Test question", "")
+
+		// Add a step that will "fail"
+		ast.BeginStep(ctx, agentcontext.StepTypeLLM, map[string]interface{}{"test": "data"})
+		// Don't complete - simulate failure
+
+		// Flush buffer (failure case)
+		testErr := fmt.Errorf("simulated error")
+		ast.FlushBuffer(ctx, agentcontext.ResumeStatusFailed, testErr)
+
+		// Verify messages were saved
+		messages, err := chatStore.GetMessages(chatID, storetypes.MessageFilter{})
+		assert.NoError(t, err)
+		assert.Len(t, messages, 1, "Should have 1 message saved")
+
+		// Verify resume records were saved
+		resumes, err := chatStore.GetResume(chatID)
+		assert.NoError(t, err)
+		assert.Len(t, resumes, 1, "Should have 1 resume record on failure")
+		assert.Equal(t, agentcontext.ResumeStatusFailed, resumes[0].Status)
+
+		// Cleanup
+		chatStore.DeleteResume(chatID)
+		chatStore.DeleteChat(chatID)
+		t.Logf("✓ Buffer flushed on failure: messages and resume records saved")
+	})
+
+	t.Run("FlushOnInterrupt", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_flush_interrupt_%s", uuid.New().String()[:8])
+		ctx := agentcontext.New(context.Background(), nil, chatID)
+		ctx.Space = plan.NewMemorySharedSpace()
+
+		// Enter stack and init buffer
+		_, _, done := agentcontext.EnterStack(ctx, ast.ID, nil)
+		defer done()
+		ast.InitBuffer(ctx)
+
+		// Ensure chat exists
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Mode:        "chat",
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+
+		// Add messages and steps
+		ctx.Buffer.AddUserInput("Test question", "")
+		ast.BeginStep(ctx, agentcontext.StepTypeLLM, nil)
+
+		// Flush buffer (interrupt case)
+		ast.FlushBuffer(ctx, agentcontext.ResumeStatusInterrupted, nil)
+
+		// Verify resume records were saved with interrupted status
+		resumes, err := chatStore.GetResume(chatID)
+		assert.NoError(t, err)
+		assert.Len(t, resumes, 1, "Should have 1 resume record on interrupt")
+		assert.Equal(t, agentcontext.ResumeStatusInterrupted, resumes[0].Status)
+
+		// Cleanup
+		chatStore.DeleteResume(chatID)
+		chatStore.DeleteChat(chatID)
+		t.Logf("✓ Buffer flushed on interrupt: resume records saved with interrupted status")
+	})
+}
+
+func TestEnsureChat(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.Get("mohe")
+	require.NoError(t, err)
+
+	// Skip if chat store not available
+	chatStore := assistant.GetChatStore()
+	if chatStore == nil {
+		t.Skip("Chat store not configured, skipping EnsureChat tests")
+	}
+
+	t.Run("CreateNewChat", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_ensure_new_%s", uuid.New().String()[:8])
+		ctx := agentcontext.New(context.Background(), nil, chatID)
+
+		// Ensure chat creates it
+		err := ast.EnsureChat(ctx)
+		assert.NoError(t, err)
+
+		// Verify chat was created
+		chat, err := chatStore.GetChat(chatID)
+		assert.NoError(t, err)
+		assert.NotNil(t, chat)
+		assert.Equal(t, chatID, chat.ChatID)
+		assert.Equal(t, ast.ID, chat.AssistantID)
+		assert.Equal(t, "active", chat.Status)
+
+		// Cleanup
+		chatStore.DeleteChat(chatID)
+		t.Logf("✓ New chat created: %s", chatID)
+	})
+
+	t.Run("SkipExistingChat", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_ensure_exist_%s", uuid.New().String()[:8])
+		ctx := agentcontext.New(context.Background(), nil, chatID)
+
+		// Create chat first
+		err := chatStore.CreateChat(&storetypes.Chat{
+			ChatID:      chatID,
+			AssistantID: ast.ID,
+			Title:       "Existing Chat",
+			Mode:        "chat",
+			Status:      "active",
+			Share:       "private",
+			CreatedAt:   time.Now(),
+			UpdatedAt:   time.Now(),
+		})
+		require.NoError(t, err)
+
+		// EnsureChat should not error
+		err = ast.EnsureChat(ctx)
+		assert.NoError(t, err)
+
+		// Verify chat still has original title
+		chat, err := chatStore.GetChat(chatID)
+		assert.NoError(t, err)
+		assert.Equal(t, "Existing Chat", chat.Title)
+
+		// Cleanup
+		chatStore.DeleteChat(chatID)
+		t.Logf("✓ Existing chat preserved")
+	})
+
+	t.Run("SkipEmptyChatID", func(t *testing.T) {
+		ctx := agentcontext.New(context.Background(), nil, "")
+
+		// Should not error with empty chat ID
+		err := ast.EnsureChat(ctx)
+		assert.NoError(t, err)
+		t.Logf("✓ Empty chat ID handled gracefully")
+	})
+
+	t.Run("CreateChatWithPermissions", func(t *testing.T) {
+		chatID := fmt.Sprintf("test_ensure_perm_%s", uuid.New().String()[:8])
+
+		// Create context with authorized info
+		ctx := agentcontext.New(context.Background(), &oauthtypes.AuthorizedInfo{
+			UserID:   "test_user_001",
+			TeamID:   "test_team_001",
+			TenantID: "test_tenant_001",
+		}, chatID)
+
+		// EnsureChat should create with permission fields
+		err := ast.EnsureChat(ctx)
+		assert.NoError(t, err)
+
+		// Verify permission fields were saved
+		chat, err := chatStore.GetChat(chatID)
+		assert.NoError(t, err)
+		assert.NotNil(t, chat)
+		assert.Equal(t, "test_user_001", chat.CreatedBy, "CreatedBy should be set")
+		assert.Equal(t, "test_user_001", chat.UpdatedBy, "UpdatedBy should be set")
+		assert.Equal(t, "test_team_001", chat.TeamID, "TeamID should be set")
+		assert.Equal(t, "test_tenant_001", chat.TenantID, "TenantID should be set")
+
+		// Cleanup
+		chatStore.DeleteChat(chatID)
+		t.Logf("✓ Chat created with permission fields: user=%s, team=%s, tenant=%s",
+			chat.CreatedBy, chat.TeamID, chat.TenantID)
+	})
+}
+
+func TestConvertBufferedTypes(t *testing.T) {
+	t.Run("ConvertBufferedMessages", func(t *testing.T) {
+		// Create buffered messages
+		buffered := []*agentcontext.BufferedMessage{
+			{
+				MessageID: "msg_001",
+				ChatID:    "chat_001",
+				RequestID: "req_001",
+				Role:      "user",
+				Type:      "user_input",
+				Props:     map[string]interface{}{"content": "Hello"},
+				Sequence:  1,
+				CreatedAt: time.Now(),
+			},
+			{
+				MessageID:   "msg_002",
+				ChatID:      "chat_001",
+				RequestID:   "req_001",
+				Role:        "assistant",
+				Type:        "text",
+				Props:       map[string]interface{}{"content": "Hi there!"},
+				BlockID:     "block_001",
+				AssistantID: "test_assistant",
+				Sequence:    2,
+				CreatedAt:   time.Now(),
+			},
+		}
+
+		// Verify structure matches store types
+		assert.Len(t, buffered, 2)
+		assert.Equal(t, "user", buffered[0].Role)
+		assert.Equal(t, "assistant", buffered[1].Role)
+		assert.Equal(t, "block_001", buffered[1].BlockID)
+		t.Logf("✓ Buffered messages have correct structure")
+	})
+
+	t.Run("ConvertBufferedSteps", func(t *testing.T) {
+		// Create buffered steps
+		buffered := []*agentcontext.BufferedStep{
+			{
+				ResumeID:      "resume_001",
+				ChatID:        "chat_001",
+				RequestID:     "req_001",
+				AssistantID:   "test_assistant",
+				StackID:       "stack_001",
+				StackDepth:    0,
+				Type:          agentcontext.StepTypeLLM,
+				Status:        agentcontext.ResumeStatusFailed,
+				Input:         map[string]interface{}{"messages": []string{"Hello"}},
+				SpaceSnapshot: map[string]interface{}{"key": "value"},
+				Error:         "Test error",
+				Sequence:      1,
+				CreatedAt:     time.Now(),
+			},
+		}
+
+		// Verify structure
+		assert.Len(t, buffered, 1)
+		assert.Equal(t, agentcontext.StepTypeLLM, buffered[0].Type)
+		assert.Equal(t, agentcontext.ResumeStatusFailed, buffered[0].Status)
+		assert.Equal(t, "Test error", buffered[0].Error)
+		assert.Equal(t, "value", buffered[0].SpaceSnapshot["key"])
+		t.Logf("✓ Buffered steps have correct structure")
+	})
+}

agent/assistant/handlers/stream.go

@@ -323,6 +323,57 @@ func (s *streamState) handleMessageEnd(data []byte) int {
 		threadID = s.ctx.Stack.ID
 	}
 
+	// Get BlockID from metadata if available
+	var blockID string
+	if s.ctx != nil {
+		if metadata := s.ctx.GetMessageMetadata(s.currentGroupID); metadata != nil {
+			blockID = metadata.BlockID
+		}
+	}
+
+	// Buffer the complete LLM message for storage
+	// Delta chunks are not stored, but we need to save the final complete content
+	// Skip if History is disabled in options
+	shouldSkipHistory := s.ctx.Stack != nil && s.ctx.Stack.Options != nil &&
+		s.ctx.Stack.Options.Skip != nil && s.ctx.Stack.Options.Skip.History
+
+	if s.ctx.Buffer != nil && len(s.buffer) > 0 && !shouldSkipHistory {
+		assistantID := ""
+		if s.ctx.Stack != nil {
+			assistantID = s.ctx.Stack.AssistantID
+		}
+
+		// Build props based on message type
+		var props map[string]interface{}
+		if msgType == message.TypeToolCall {
+			// For tool calls, try to parse the accumulated buffer as JSON
+			var toolCallData interface{}
+			if err := jsoniter.Unmarshal(s.buffer, &toolCallData); err == nil {
+				props = map[string]interface{}{
+					"calls": toolCallData,
+				}
+			} else {
+				props = map[string]interface{}{
+					"content": string(s.buffer),
+				}
+			}
+		} else {
+			// For text/thinking, content is the accumulated text
+			props = map[string]interface{}{
+				"content": string(s.buffer),
+			}
+		}
+
+		s.ctx.Buffer.AddAssistantMessage(
+			msgType,
+			props,
+			blockID,
+			threadID,
+			assistantID,
+			nil,
+		)
+	}
+
 	// Build EventMessageEndData with complete content
 	endData := message.EventMessageEndData{
 		MessageID:  s.currentGroupID, // Use the message ID

agent/context/buffer.go

@@ -0,0 +1,359 @@
+package context
+
+import (
+	"sync"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+// =============================================================================
+// Chat Buffer - Buffers messages and steps during execution for batch saving
+// =============================================================================
+
+// ChatBuffer buffers messages and resume steps during agent execution
+// All data is held in memory and batch-written at the end of Stream()
+type ChatBuffer struct {
+	// Identity
+	chatID      string
+	requestID   string
+	assistantID string
+
+	// Message buffer
+	messages    []*BufferedMessage
+	msgSequence int
+
+	// Step buffer (for Resume)
+	steps        []*BufferedStep
+	currentStep  *BufferedStep
+	stepSequence int
+
+	// Space snapshot (captured when step starts, for recovery)
+	spaceSnapshot map[string]interface{}
+
+	mu sync.Mutex
+}
+
+// BufferedMessage represents a message waiting to be saved
+type BufferedMessage struct {
+	MessageID   string                 `json:"message_id"`
+	ChatID      string                 `json:"chat_id"`
+	RequestID   string                 `json:"request_id,omitempty"`
+	Role        string                 `json:"role"` // "user" or "assistant"
+	Type        string                 `json:"type"` // "text", "image", "loading", "tool_call", "retrieval", etc.
+	Props       map[string]interface{} `json:"props"`
+	BlockID     string                 `json:"block_id,omitempty"`
+	ThreadID    string                 `json:"thread_id,omitempty"`
+	AssistantID string                 `json:"assistant_id,omitempty"`
+	Sequence    int                    `json:"sequence"`
+	Metadata    map[string]interface{} `json:"metadata,omitempty"`
+	CreatedAt   time.Time              `json:"created_at"`
+}
+
+// BufferedStep represents an execution step waiting to be saved (for Resume)
+// Only saved when request is interrupted or failed
+type BufferedStep struct {
+	ResumeID      string                 `json:"resume_id"`
+	ChatID        string                 `json:"chat_id"`
+	RequestID     string                 `json:"request_id"`
+	AssistantID   string                 `json:"assistant_id"`
+	StackID       string                 `json:"stack_id"`
+	StackParentID string                 `json:"stack_parent_id,omitempty"`
+	StackDepth    int                    `json:"stack_depth"`
+	Type          string                 `json:"type"`   // "input", "hook_create", "llm", "tool", "hook_next", "delegate"
+	Status        string                 `json:"status"` // "running", "completed", "failed", "interrupted"
+	Input         map[string]interface{} `json:"input,omitempty"`
+	Output        map[string]interface{} `json:"output,omitempty"`
+	SpaceSnapshot map[string]interface{} `json:"space_snapshot,omitempty"`
+	Error         string                 `json:"error,omitempty"`
+	Sequence      int                    `json:"sequence"`
+	Metadata      map[string]interface{} `json:"metadata,omitempty"`
+	CreatedAt     time.Time              `json:"created_at"`
+}
+
+// Step status constants (internal use only, not stored in database)
+const (
+	StepStatusRunning   = "running"
+	StepStatusCompleted = "completed"
+)
+
+// Step type constants
+const (
+	StepTypeInput      = "input"
+	StepTypeHookCreate = "hook_create"
+	StepTypeLLM        = "llm"
+	StepTypeTool       = "tool"
+	StepTypeHookNext   = "hook_next"
+	StepTypeDelegate   = "delegate"
+)
+
+// Resume status constants (for database storage)
+const (
+	ResumeStatusFailed      = "failed"
+	ResumeStatusInterrupted = "interrupted"
+)
+
+// NewChatBuffer creates a new chat buffer
+func NewChatBuffer(chatID, requestID, assistantID string) *ChatBuffer {
+	return &ChatBuffer{
+		chatID:      chatID,
+		requestID:   requestID,
+		assistantID: assistantID,
+		messages:    make([]*BufferedMessage, 0),
+		steps:       make([]*BufferedStep, 0),
+	}
+}
+
+// =============================================================================
+// Message Buffer Methods
+// =============================================================================
+
+// AddMessage adds a message to the buffer
+func (b *ChatBuffer) AddMessage(msg *BufferedMessage) {
+	if msg == nil {
+		return
+	}
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	// Auto-generate IDs if not provided
+	if msg.MessageID == "" {
+		msg.MessageID = uuid.New().String()
+	}
+	if msg.ChatID == "" {
+		msg.ChatID = b.chatID
+	}
+	if msg.RequestID == "" {
+		msg.RequestID = b.requestID
+	}
+	if msg.CreatedAt.IsZero() {
+		msg.CreatedAt = time.Now()
+	}
+
+	// Auto-increment sequence
+	b.msgSequence++
+	msg.Sequence = b.msgSequence
+
+	b.messages = append(b.messages, msg)
+}
+
+// AddUserInput adds user input message to the buffer
+func (b *ChatBuffer) AddUserInput(content interface{}, name string) {
+	props := map[string]interface{}{
+		"content": content,
+		"role":    "user",
+	}
+	if name != "" {
+		props["name"] = name
+	}
+
+	b.AddMessage(&BufferedMessage{
+		Role:  "user",
+		Type:  "user_input",
+		Props: props,
+	})
+}
+
+// AddAssistantMessage adds an assistant message to the buffer
+// This is called by ctx.Send() to buffer messages for batch saving
+func (b *ChatBuffer) AddAssistantMessage(msgType string, props map[string]interface{}, blockID, threadID, assistantID string, metadata map[string]interface{}) {
+	// Skip event type messages (transient, not stored)
+	if msgType == "event" {
+		return
+	}
+
+	b.AddMessage(&BufferedMessage{
+		Role:        "assistant",
+		Type:        msgType,
+		Props:       props,
+		BlockID:     blockID,
+		ThreadID:    threadID,
+		AssistantID: assistantID,
+		Metadata:    metadata,
+	})
+}
+
+// GetMessages returns all buffered messages
+func (b *ChatBuffer) GetMessages() []*BufferedMessage {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	result := make([]*BufferedMessage, len(b.messages))
+	copy(result, b.messages)
+	return result
+}
+
+// GetMessageCount returns the number of buffered messages
+func (b *ChatBuffer) GetMessageCount() int {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return len(b.messages)
+}
+
+// =============================================================================
+// Step Buffer Methods (for Resume)
+// =============================================================================
+
+// BeginStep starts tracking a new execution step
+// Returns the step for further updates
+func (b *ChatBuffer) BeginStep(stepType string, input map[string]interface{}, stack *Stack) *BufferedStep {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	b.stepSequence++
+
+	step := &BufferedStep{
+		ResumeID:    uuid.New().String(),
+		ChatID:      b.chatID,
+		RequestID:   b.requestID,
+		AssistantID: b.assistantID,
+		Type:        stepType,
+		Status:      StepStatusRunning,
+		Input:       input,
+		Sequence:    b.stepSequence,
+		CreatedAt:   time.Now(),
+	}
+
+	// Set stack information if available
+	if stack != nil {
+		step.StackID = stack.ID
+		step.StackParentID = stack.ParentID
+		step.StackDepth = stack.Depth
+	}
+
+	// Capture current space snapshot
+	if b.spaceSnapshot != nil {
+		step.SpaceSnapshot = copyMap(b.spaceSnapshot)
+	}
+
+	b.steps = append(b.steps, step)
+	b.currentStep = step
+
+	return step
+}
+
+// CompleteStep marks the current step as completed
+func (b *ChatBuffer) CompleteStep(output map[string]interface{}) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if b.currentStep != nil {
+		b.currentStep.Output = output
+		b.currentStep.Status = StepStatusCompleted
+		b.currentStep = nil
+	}
+}
+
+// FailCurrentStep marks the current step as failed or interrupted
+func (b *ChatBuffer) FailCurrentStep(status string, err error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if b.currentStep != nil && b.currentStep.Status == StepStatusRunning {
+		b.currentStep.Status = status
+		if err != nil {
+			b.currentStep.Error = err.Error()
+		}
+	}
+}
+
+// GetCurrentStep returns the current running step
+func (b *ChatBuffer) GetCurrentStep() *BufferedStep {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.currentStep
+}
+
+// GetStepsForResume returns steps that need to be saved for resume
+// Only returns steps with failed or interrupted status
+func (b *ChatBuffer) GetStepsForResume(finalStatus string) []*BufferedStep {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	// If completed successfully, no steps need to be saved
+	if finalStatus == StepStatusCompleted {
+		return nil
+	}
+
+	// Mark current running step with final status
+	if b.currentStep != nil && b.currentStep.Status == StepStatusRunning {
+		b.currentStep.Status = finalStatus
+	}
+
+	// Return all steps (they will all have the context for recovery)
+	result := make([]*BufferedStep, len(b.steps))
+	copy(result, b.steps)
+	return result
+}
+
+// GetAllSteps returns all buffered steps (for debugging/testing)
+func (b *ChatBuffer) GetAllSteps() []*BufferedStep {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	result := make([]*BufferedStep, len(b.steps))
+	copy(result, b.steps)
+	return result
+}
+
+// =============================================================================
+// Space Snapshot Methods
+// =============================================================================
+
+// SetSpaceSnapshot sets the space snapshot for recovery
+// Should be called when space data changes
+func (b *ChatBuffer) SetSpaceSnapshot(snapshot map[string]interface{}) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.spaceSnapshot = copyMap(snapshot)
+}
+
+// GetSpaceSnapshot returns the current space snapshot
+func (b *ChatBuffer) GetSpaceSnapshot() map[string]interface{} {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return copyMap(b.spaceSnapshot)
+}
+
+// =============================================================================
+// Identity Methods
+// =============================================================================
+
+// ChatID returns the chat ID
+func (b *ChatBuffer) ChatID() string {
+	return b.chatID
+}
+
+// RequestID returns the request ID
+func (b *ChatBuffer) RequestID() string {
+	return b.requestID
+}
+
+// AssistantID returns the assistant ID
+func (b *ChatBuffer) AssistantID() string {
+	return b.assistantID
+}
+
+// SetAssistantID updates the assistant ID (for A2A calls)
+func (b *ChatBuffer) SetAssistantID(assistantID string) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.assistantID = assistantID
+}
+
+// =============================================================================
+// Helper Functions
+// =============================================================================
+
+// copyMap creates a shallow copy of a map
+func copyMap(src map[string]interface{}) map[string]interface{} {
+	if src == nil {
+		return nil
+	}
+	dst := make(map[string]interface{}, len(src))
+	for k, v := range src {
+		dst[k] = v
+	}
+	return dst
+}

agent/context/context.go

@@ -387,3 +387,93 @@ func (ctx *Context) getMessageMetadata(messageID string) *MessageMetadata {
 	}
 	return ctx.messageMetadata.getMessage(messageID)
 }
+
+// GetMessageMetadata returns metadata for a message (public version)
+func (ctx *Context) GetMessageMetadata(messageID string) *MessageMetadata {
+	return ctx.getMessageMetadata(messageID)
+}
+
+// =============================================================================
+// Chat Buffer Methods
+// =============================================================================
+
+// InitBuffer initializes the chat buffer for this context
+// Should be called at the start of Stream() to begin buffering messages and steps
+func (ctx *Context) InitBuffer(assistantID string) *ChatBuffer {
+	ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID)
+	return ctx.Buffer
+}
+
+// HasBuffer returns true if the buffer is initialized
+func (ctx *Context) HasBuffer() bool {
+	return ctx.Buffer != nil
+}
+
+// BufferUserInput adds user input to the buffer
+// Should be called at the start of Stream() to buffer the user's input message
+func (ctx *Context) BufferUserInput(messages []Message) {
+	if ctx.Buffer == nil {
+		return
+	}
+
+	for _, msg := range messages {
+		if msg.Role == RoleUser {
+			// Get name if available
+			var name string
+			if msg.Name != nil {
+				name = *msg.Name
+			}
+			ctx.Buffer.AddUserInput(msg.Content, name)
+		}
+	}
+}
+
+// BufferAssistantMessage adds an assistant message to the buffer
+// Called by ctx.Send() to buffer messages for batch saving
+func (ctx *Context) BufferAssistantMessage(msgType string, props map[string]interface{}, blockID, threadID string, metadata map[string]interface{}) {
+	if ctx.Buffer == nil {
+		return
+	}
+
+	ctx.Buffer.AddAssistantMessage(msgType, props, blockID, threadID, ctx.AssistantID, metadata)
+}
+
+// BeginStep starts tracking a new execution step
+// Returns the step for further updates
+func (ctx *Context) BeginStep(stepType string, input map[string]interface{}) *BufferedStep {
+	if ctx.Buffer == nil {
+		return nil
+	}
+
+	// Update space snapshot before starting step
+	if ctx.Space != nil {
+		ctx.Buffer.SetSpaceSnapshot(ctx.Space.Snapshot())
+	}
+
+	return ctx.Buffer.BeginStep(stepType, input, ctx.Stack)
+}
+
+// CompleteStep marks the current step as completed
+func (ctx *Context) CompleteStep(output map[string]interface{}) {
+	if ctx.Buffer == nil {
+		return
+	}
+	ctx.Buffer.CompleteStep(output)
+}
+
+// FailCurrentStep marks the current step as failed or interrupted
+func (ctx *Context) FailCurrentStep(status string, err error) {
+	if ctx.Buffer == nil {
+		return
+	}
+	ctx.Buffer.FailCurrentStep(status, err)
+}
+
+// shouldSkipHistory checks if history saving should be skipped
+// Returns true if Skip.History is set in the current stack options
+func (ctx *Context) shouldSkipHistory() bool {
+	if ctx.Stack == nil || ctx.Stack.Options == nil || ctx.Stack.Options.Skip == nil {
+		return false
+	}
+	return ctx.Stack.Options.Skip.History
+}

agent/context/output.go

@@ -146,6 +146,25 @@ func (ctx *Context) Send(msg *message.Message) error {
 		return err
 	}
 
+	// === Buffer message for batch saving (non-delta, non-event messages only) ===
+	// Delta messages are streaming chunks; only final content should be saved
+	// Event messages are transient lifecycle signals, not stored
+	// Skip if History is disabled in options
+	if !msg.Delta && !isEventMessage && ctx.Buffer != nil && !ctx.shouldSkipHistory() {
+		assistantID := ""
+		if ctx.Stack != nil {
+			assistantID = ctx.Stack.AssistantID
+		}
+		ctx.Buffer.AddAssistantMessage(
+			msg.Type,
+			msg.Props,
+			msg.BlockID,
+			msg.ThreadID,
+			assistantID,
+			nil, // metadata can be added if needed
+		)
+	}
+
 	// === 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)

agent/context/types.go

@@ -231,6 +231,9 @@ type Context struct {
 	Writer      Writer               `json:"-"`  // Writer, 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)
 
+	// Chat buffer for batch saving messages and resume steps
+	Buffer *ChatBuffer `json:"-"` // Chat buffer for batch saving at end of Stream()
+
 	// Internal
 	trace           traceTypes.Manager    `json:"-"` // Trace manager, lazy initialized on first access
 	messageMetadata *messageMetadataStore `json:"-"` // Thread-safe message metadata store for delta operations

agent/store/types/types.go

@@ -38,6 +38,12 @@ type Chat struct {
 	Metadata      map[string]interface{} `json:"metadata,omitempty"`
 	CreatedAt     time.Time              `json:"created_at"`
 	UpdatedAt     time.Time              `json:"updated_at"`
+
+	// Permission fields (managed by Yao framework when permission: true)
+	CreatedBy string `json:"__yao_created_by,omitempty"` // User ID who created the record
+	UpdatedBy string `json:"__yao_updated_by,omitempty"` // User ID who last updated
+	TeamID    string `json:"__yao_team_id,omitempty"`    // Team ID for team-level access
+	TenantID  string `json:"__yao_tenant_id,omitempty"`  // Tenant ID for multi-tenancy
 }
 
 // ChatFilter for listing chats

agent/store/xun/chat.go

@@ -80,6 +80,20 @@ func (store *Xun) CreateChat(chat *types.Chat) error {
 		data["metadata"] = metadataJSON
 	}
 
+	// Handle permission fields (Yao framework permission: true)
+	if chat.CreatedBy != "" {
+		data["__yao_created_by"] = chat.CreatedBy
+	}
+	if chat.UpdatedBy != "" {
+		data["__yao_updated_by"] = chat.UpdatedBy
+	}
+	if chat.TeamID != "" {
+		data["__yao_team_id"] = chat.TeamID
+	}
+	if chat.TenantID != "" {
+		data["__yao_tenant_id"] = chat.TenantID
+	}
+
 	// Insert
 	return store.newQueryChat().Insert(data)
 }
@@ -349,6 +363,12 @@ func (store *Xun) rowToChat(data map[string]interface{}) (*types.Chat, error) {
 		}
 	}
 
+	// Handle permission fields
+	chat.CreatedBy = getString(data, "__yao_created_by")
+	chat.UpdatedBy = getString(data, "__yao_updated_by")
+	chat.TeamID = getString(data, "__yao_team_id")
+	chat.TenantID = getString(data, "__yao_tenant_id")
+
 	return chat, nil
 }