Enhance tracing and message handling in Assistant context

- Added error handling and logging for completion node creation in the traceAgentCompletion method, improving traceability of failures.
- Introduced a new replaceMethod in the context to allow message content replacement, enhancing message management capabilities.
- Updated documentation for sendMethod and replaceMethod to clarify usage and return values, improving developer experience.

agent/assistant/trace.go

@@ -3,6 +3,7 @@ package assistant
 import (
 	"fmt"
 
+	"github.com/yaoapp/kun/log"
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/i18n"
 	"github.com/yaoapp/yao/trace/types"
@@ -98,7 +99,7 @@ func (ast *Assistant) traceAgentCompletion(ctx *context.Context, createResponse
 	}
 
 	// Create a dedicated completion node
-	completionNode, _ := trace.Add(
+	completionNode, err := trace.Add(
 		input,
 		types.TraceNodeOption{
 			Label:       i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.completion.label"), // "Agent Completion"
@@ -107,6 +108,10 @@ func (ast *Assistant) traceAgentCompletion(ctx *context.Context, createResponse
 			Description: i18n.Tr(ast.ID, ctx.Locale, "assistant.agent.completion.description"), // "Final output from assistant"
 		},
 	)
+	if err != nil {
+		log.Trace("[TRACE] Failed to create completion node: %v", err)
+		return
+	}
 
 	// Immediately mark it as complete with the final response
 	if completionNode != nil {

agent/context/JSAPI.md

@@ -0,0 +1,621 @@
+# Context JavaScript API Documentation
+
+## Overview
+
+The Context JavaScript API provides a comprehensive interface for interacting with the Yao Agent system from JavaScript/TypeScript hooks (Create, Next, Done). The Context object exposes agent state, configuration, messaging capabilities, trace operations, and MCP (Model Context Protocol) integrations.
+
+## Context Object
+
+The Context object is automatically passed to hook functions and provides access to the agent's execution environment.
+
+### Basic Properties
+
+```typescript
+interface Context {
+  // Identifiers
+  chat_id: string; // Current chat session ID
+  assistant_id: string; // Assistant identifier
+
+  // Configuration
+  connector: string; // LLM connector name
+  search?: string; // Search engine configuration
+  locale: string; // User locale (e.g., "en", "zh-cn")
+  theme: string; // UI theme preference
+  accept: string; // Output format ("openai", "cui", etc.)
+  route: string; // Request route path
+  referer: string; // Request referer
+
+  // Retry Configuration
+  retry: boolean; // Whether retry is enabled
+  retry_times: number; // Number of retry attempts
+
+  // Client Information
+  client: {
+    type: string; // Client type
+    user_agent: string; // User agent string
+    ip: string; // Client IP address
+  };
+
+  // Dynamic Data
+  args?: any[]; // Additional arguments
+  metadata?: Record<string, any>; // Custom metadata
+  authorized?: Record<string, any>; // Authorization data
+}
+```
+
+## Methods
+
+### Send Messages
+
+#### `ctx.Send(message): string`
+
+Sends a message to the client and automatically flushes the output.
+
+**Parameters:**
+
+- `message`: Message object or string
+
+**Returns:**
+
+- `string`: The message ID (auto-generated if not provided in the message object)
+
+**Message Object Structure:**
+
+```typescript
+interface Message {
+  type: string; // Message type: "text", "tool", "image", etc.
+  props: Record<string, any>; // Message properties
+  message_id?: string; // Optional message ID (auto-generated if omitted)
+}
+```
+
+**Examples:**
+
+```javascript
+// Send text message (object format) and capture message ID
+const messageId = ctx.Send({
+  type: "text",
+  props: { content: "Hello, World!" },
+});
+console.log("Sent message:", messageId);
+
+// Send text message (shorthand)
+const textId = ctx.Send("Hello, World!");
+
+// Send tool message with custom ID
+const toolId = ctx.Send({
+  type: "tool",
+  message_id: "custom-tool-msg-1",
+  props: {
+    name: "calculator",
+    result: { sum: 42 },
+  },
+});
+
+// Send image message
+const imageId = ctx.Send({
+  type: "image",
+  props: {
+    url: "https://example.com/image.png",
+    alt: "Example Image",
+  },
+});
+```
+
+**Notes:**
+
+- Message ID is automatically generated if not provided
+- Returns the message ID for reference in subsequent operations
+- Output is automatically flushed after sending
+- Throws exception on failure
+
+#### `ctx.Replace(messageId, message): string`
+
+Replaces an existing message with new content. This is useful for updating progress messages or correcting previously sent information.
+
+**Parameters:**
+
+- `messageId`: String - The ID of the message to replace
+- `message`: Message object or string - The new message content
+
+**Returns:**
+
+- `string`: The message ID (same as the provided messageId)
+
+**Examples:**
+
+```javascript
+// Send initial message
+const msgId = ctx.Send("Processing...");
+
+// Later, replace with updated content
+ctx.Replace(msgId, "Processing complete!");
+
+// Replace with complex message
+ctx.Replace(msgId, {
+  type: "text",
+  props: {
+    content: "Task finished",
+    status: "success",
+  },
+});
+
+// Replace with shorthand text
+ctx.Replace(msgId, "Updated text content");
+```
+
+**Use Cases:**
+
+```javascript
+// Progress updates
+const progressId = ctx.Send("Step 1/3: Starting...");
+// ... do work ...
+ctx.Replace(progressId, "Step 2/3: Processing...");
+// ... do more work ...
+ctx.Replace(progressId, "Step 3/3: Finalizing...");
+// ... finish ...
+ctx.Replace(progressId, "Complete! ✓");
+
+// Error correction
+const msgId = ctx.Send("Found 5 results");
+// Oops, counted wrong
+ctx.Replace(msgId, "Found 8 results");
+```
+
+**Notes:**
+
+- The message must exist (must have been sent previously)
+- Replaces the entire message content, not just specific fields
+- Output is automatically flushed after replacing
+- Throws exception on failure
+
+### Resource Cleanup
+
+#### `ctx.Release()`
+
+Manually releases Context resources. This is optional as cleanup happens automatically via garbage collection.
+
+**Example:**
+
+```javascript
+try {
+  // Use context
+  ctx.Send("Processing...");
+} finally {
+  ctx.Release(); // Manual cleanup
+}
+```
+
+## Trace API
+
+The `ctx.Trace` object provides comprehensive tracing capabilities for debugging and monitoring agent execution.
+
+### Node Operations
+
+#### `ctx.Trace.Add(input, options)`
+
+Creates a new trace node (sequential step).
+
+**Parameters:**
+
+- `input`: Input data for the node
+- `options`: Node configuration object
+
+**Options Structure:**
+
+```typescript
+interface TraceNodeOption {
+  label: string; // Display label
+  type: string; // Node type identifier
+  icon: string; // Icon identifier
+  description: string; // Node description
+  metadata?: Record<string, any>; // Additional metadata
+}
+```
+
+**Example:**
+
+```javascript
+const node = ctx.Trace.Add(
+  { query: "What is AI?" },
+  {
+    label: "Search Query",
+    type: "search",
+    icon: "search",
+    description: "Searching for AI information",
+  }
+);
+```
+
+#### `ctx.Trace.Parallel(inputs)`
+
+Creates multiple parallel trace nodes for concurrent operations.
+
+**Parameters:**
+
+- `inputs`: Array of parallel input objects
+
+**Input Structure:**
+
+```typescript
+interface ParallelInput {
+  input: any; // Input data
+  option: TraceNodeOption; // Node configuration
+}
+```
+
+**Example:**
+
+```javascript
+const nodes = ctx.Trace.Parallel([
+  {
+    input: { url: "https://api1.com" },
+    option: {
+      label: "API Call 1",
+      type: "api",
+      icon: "cloud",
+      description: "Fetching from API 1",
+    },
+  },
+  {
+    input: { url: "https://api2.com" },
+    option: {
+      label: "API Call 2",
+      type: "api",
+      icon: "cloud",
+      description: "Fetching from API 2",
+    },
+  },
+]);
+```
+
+### Logging Methods
+
+Add log entries to the current trace node:
+
+```javascript
+// Information logs
+ctx.Trace.Info("Processing started", { step: 1 });
+
+// Debug logs
+ctx.Trace.Debug("Variable value", { value: 42 });
+
+// Warning logs
+ctx.Trace.Warn("Deprecated feature used", { feature: "old_api" });
+
+// Error logs
+ctx.Trace.Error("Operation failed", { error: "timeout" });
+```
+
+### Node Status Operations
+
+#### `node.SetOutput(output)`
+
+Sets the output data for a node.
+
+```javascript
+const node = ctx.Trace.Add({ query: "search" }, options);
+node.SetOutput({ results: [...] });
+```
+
+#### `node.SetMetadata(key, value)`
+
+Sets metadata for a node.
+
+```javascript
+node.SetMetadata("duration", 1500);
+node.SetMetadata("cache_hit", true);
+```
+
+#### `node.Complete(output?)`
+
+Marks a node as completed (optionally with output).
+
+```javascript
+node.Complete({ status: "success", data: [...] });
+```
+
+#### `node.Fail(error)`
+
+Marks a node as failed with an error.
+
+```javascript
+try {
+  // Operation
+} catch (error) {
+  node.Fail(error);
+}
+```
+
+### Query Operations
+
+#### `ctx.Trace.GetRootNode()`
+
+Returns the root node of the trace tree.
+
+```javascript
+const root = ctx.Trace.GetRootNode();
+console.log(root.id, root.label);
+```
+
+#### `ctx.Trace.GetNode(id)`
+
+Retrieves a specific node by ID.
+
+```javascript
+const node = ctx.Trace.GetNode("node-123");
+```
+
+#### `ctx.Trace.GetCurrentNodes()`
+
+Returns the current active nodes (may be multiple if in parallel state).
+
+```javascript
+const currentNodes = ctx.Trace.GetCurrentNodes();
+```
+
+### Memory Space Operations
+
+#### `ctx.Trace.CreateSpace(option)`
+
+Creates a memory space for storing key-value data.
+
+```javascript
+const space = ctx.Trace.CreateSpace({
+  label: "Context Memory",
+  type: "context",
+  icon: "database",
+  description: "Stores conversation context",
+});
+```
+
+#### `ctx.Trace.GetSpace(id)`
+
+Retrieves a memory space by ID.
+
+```javascript
+const space = ctx.Trace.GetSpace("context");
+```
+
+#### `ctx.Trace.HasSpace(id)`
+
+Checks if a memory space exists.
+
+```javascript
+if (ctx.Trace.HasSpace("context")) {
+  // Space exists
+}
+```
+
+#### `ctx.Trace.DeleteSpace(id)`
+
+Deletes a memory space.
+
+```javascript
+ctx.Trace.DeleteSpace("temp_storage");
+```
+
+#### `ctx.Trace.ListSpaces()`
+
+Lists all memory spaces.
+
+```javascript
+const spaces = ctx.Trace.ListSpaces();
+spaces.forEach((space) => {
+  console.log(space.id, space.label);
+});
+```
+
+## MCP API
+
+The `ctx.MCP` object provides access to Model Context Protocol operations for interacting with external tools, resources, and prompts.
+
+### Resource Operations
+
+#### `ctx.MCP.ListResources(client)`
+
+Lists available resources from an MCP client.
+
+```javascript
+const resources = ctx.MCP.ListResources("filesystem");
+```
+
+#### `ctx.MCP.ReadResource(client, uri)`
+
+Reads a specific resource.
+
+```javascript
+const content = ctx.MCP.ReadResource("filesystem", "file:///path/to/file.txt");
+```
+
+### Tool Operations
+
+#### `ctx.MCP.ListTools(client)`
+
+Lists available tools from an MCP client.
+
+```javascript
+const tools = ctx.MCP.ListTools("toolkit");
+```
+
+#### `ctx.MCP.CallTool(client, name, args)`
+
+Calls a single tool.
+
+```javascript
+const result = ctx.MCP.CallTool("calculator", "add", {
+  a: 10,
+  b: 32,
+});
+```
+
+#### `ctx.MCP.CallTools(client, calls)`
+
+Calls multiple tools sequentially.
+
+```javascript
+const results = ctx.MCP.CallTools("toolkit", [
+  { name: "tool1", args: { param: "value1" } },
+  { name: "tool2", args: { param: "value2" } },
+]);
+```
+
+#### `ctx.MCP.CallToolsParallel(client, calls)`
+
+Calls multiple tools in parallel.
+
+```javascript
+const results = ctx.MCP.CallToolsParallel("toolkit", [
+  { name: "api1", args: { endpoint: "/users" } },
+  { name: "api2", args: { endpoint: "/posts" } },
+]);
+```
+
+### Prompt Operations
+
+#### `ctx.MCP.ListPrompts(client)`
+
+Lists available prompts from an MCP client.
+
+```javascript
+const prompts = ctx.MCP.ListPrompts("prompt_library");
+```
+
+#### `ctx.MCP.GetPrompt(client, name, args?)`
+
+Retrieves a specific prompt.
+
+```javascript
+const prompt = ctx.MCP.GetPrompt("prompt_library", "code_review", {
+  language: "javascript",
+});
+```
+
+### Sample Operations
+
+#### `ctx.MCP.CreateSample(client, uri, sample)`
+
+Creates a sample for a resource.
+
+```javascript
+ctx.MCP.CreateSample("filesystem", "file:///examples", {
+  name: "example1",
+  content: "Sample content",
+});
+```
+
+## Complete Example
+
+Here's a comprehensive example using various Context API features:
+
+```javascript
+/**
+ * Next Hook - Process LLM response and enhance with tools
+ */
+function Next(ctx, messages, completion, tools) {
+  try {
+    // Create trace node for custom processing
+    const processNode = ctx.Trace.Add(
+      { completion, tools },
+      {
+        label: "Custom Processing",
+        type: "custom",
+        icon: "settings",
+        description: "Enhancing response with external data",
+      }
+    );
+
+    // Log processing start
+    ctx.Trace.Info("Starting custom processing", {
+      tool_count: tools?.length || 0,
+    });
+
+    // Send progress message and capture message ID
+    const progressId = ctx.Send("Searching for articles...");
+
+    // Call MCP tool for additional data
+    const searchResults = ctx.MCP.CallTool("search_engine", "search", {
+      query: "latest AI news",
+      limit: 5,
+    });
+
+    // Update trace with results
+    processNode.SetMetadata("search_results_count", searchResults.length);
+
+    // Update the progress message with results
+    ctx.Replace(progressId, `Found ${searchResults.length} relevant articles.`);
+
+    // Log the message ID for tracking
+    ctx.Trace.Debug("Updated progress message", { message_id: progressId });
+
+    // Process and format response
+    const enhancedResponse = {
+      text: completion.content,
+      sources: searchResults,
+      timestamp: Date.now(),
+    };
+
+    // Mark node as complete
+    processNode.Complete(enhancedResponse);
+
+    // Return enhanced response
+    return {
+      data: enhancedResponse,
+      done: true,
+    };
+  } catch (error) {
+    ctx.Trace.Error("Processing failed", { error: error.message });
+    throw error;
+  } finally {
+    // Optional: Manual cleanup
+    ctx.Release();
+  }
+}
+```
+
+## Best Practices
+
+1. **Error Handling**: Always wrap Context operations in try-catch blocks
+2. **Resource Cleanup**: Use try-finally pattern for manual cleanup if needed
+3. **Trace Organization**: Create meaningful trace nodes with descriptive labels
+4. **Logging Levels**: Use appropriate log levels (Debug for development, Info for progress, Error for failures)
+5. **Message IDs**: Let the system auto-generate message IDs unless you need specific tracking
+6. **Parallel Operations**: Use `Trace.Parallel()` for concurrent operations to maintain trace clarity
+7. **Memory Spaces**: Use memory spaces for persistent data across agent calls
+
+## Error Handling
+
+All Context methods throw exceptions on failure. Always handle errors appropriately:
+
+```javascript
+try {
+  ctx.Send(message);
+} catch (error) {
+  ctx.Trace.Error("Failed to send message", { error: error.message });
+  throw error;
+}
+```
+
+## TypeScript Support
+
+For TypeScript projects, the Context types are automatically inferred. You can also import explicit types:
+
+```typescript
+import { Context, Message, TraceNodeOption } from "@yaoapps/types";
+
+function Next(
+  ctx: Context,
+  messages: Message[],
+  completion: any,
+  tools: any[]
+): any {
+  // Your code with full type checking
+}
+```
+
+## See Also
+
+- [Agent Hooks Documentation](../hooks/README.md)
+- [MCP Protocol Specification](../mcp/README.md)
+- [Trace System Documentation](../../trace/README.md)
+- [Message Format Specification](../message/README.md)

agent/context/jsapi.go

@@ -3,6 +3,7 @@ package context
 import (
 	"github.com/yaoapp/gou/runtime/v8/bridge"
 	"github.com/yaoapp/yao/agent/output"
+	"github.com/yaoapp/yao/agent/output/message"
 	traceJsapi "github.com/yaoapp/yao/trace/jsapi"
 	"rogchap.com/v8go"
 )
@@ -50,6 +51,7 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) {
 
 	// Set methods
 	jsObject.Set("Send", ctx.sendMethod(v8ctx.Isolate()))
+	jsObject.Set("Replace", ctx.replaceMethod(v8ctx.Isolate()))
 
 	// Set MCP object
 	jsObject.Set("MCP", ctx.newMCPObject(v8ctx.Isolate()))
@@ -134,17 +136,24 @@ func (ctx *Context) objectRelease(iso *v8go.Isolate, goValueID string) *v8go.Fun
 		// Get the context object (this)
 		thisObj, err := info.This().AsObject()
 		if err == nil {
-			// Release Trace object if it has __release method
+			// NOTE: We do NOT automatically release Trace object here
-			if traceVal, err := thisObj.Get("Trace"); err == nil && !traceVal.IsNullOrUndefined() {
+			//
-				if traceObj, err := traceVal.AsObject(); err == nil {
+			// Rationale:
-					if releaseFunc, err := traceObj.Get("__release"); err == nil && releaseFunc.IsFunction() {
+			// 1. Each Hook execution creates a new V8 script context (scriptCtx)
-						// Call Trace.__release() to cleanup trace resources
+			// 2. The agent Context (ctx) is passed to the Hook as a parameter
-						if releaseFn, err := releaseFunc.AsFunction(); err == nil {
+			// 3. When scriptCtx.Close() is called (via defer), V8 cleanup triggers ctx.__release()
-							releaseFn.Call(traceObj.Value) // Ignore errors in cleanup
+			// 4. If we release Trace here, it gets released after EVERY Hook execution
-						}
+			// 5. This causes "context canceled" errors in subsequent operations
-					}
+			//
-				}
+			// Trace lifecycle:
-			}
+			// - Trace is created when agent.Stream() starts (in Context.Trace())
+			// - Trace should persist across ALL Hook executions (Create, Next, Done)
+			// - Trace is released when agent Context.Release() is called (after agent.Stream() completes)
+			//
+			// Memory management:
+			// - If JS code explicitly calls trace.Release(), it will work (trace/jsapi/trace.go:traceGoRelease)
+			// - If not explicitly called, Context.Release() will clean it up (context/context.go:Release)
+			// - This is the correct lifecycle: one Context -> one Trace -> multiple Hook executions
 
 			// Release Context Go object from bridge registry
 			if thisObj.InternalFieldCount() > 0 {
@@ -190,9 +199,10 @@ func (ctx *Context) createTraceObject(v8ctx *v8go.Context) *v8go.Value {
 }
 
 // sendMethod implements ctx.Send(message)
-// Usage: ctx.Send({ type: "text", props: { content: "Hello" } })
+// Usage: const messageId = ctx.Send({ type: "text", props: { content: "Hello" } })
-// Usage: ctx.Send("Hello") // shorthand for text message
+// Usage: const messageId = ctx.Send("Hello") // shorthand for text message
 // Automatically generates ID and flushes output
+// Returns: message_id (string)
 func (ctx *Context) sendMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
 	return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value {
 		v8ctx := info.Context()
@@ -227,7 +237,66 @@ func (ctx *Context) sendMethod(iso *v8go.Isolate) *v8go.FunctionTemplate {
 			return bridge.JsException(v8ctx, "Flush failed: "+err.Error())
 		}
 
-		return v8go.Undefined(iso)
+		// Return the message ID
+		messageID, err := v8go.NewValue(iso, msg.MessageID)
+		if err != nil {
+			return bridge.JsException(v8ctx, "Failed to create return value: "+err.Error())
+		}
+		return messageID
+	})
+}
+
+// replaceMethod implements ctx.Replace(messageId, message)
+// Usage: ctx.Replace(messageId, { type: "text", props: { content: "Updated content" } })
+// Replaces the entire message content with the specified message_id
+// Automatically flushes output
+// Returns: message_id (string)
+func (ctx *Context) replaceMethod(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) < 2 {
+			return bridge.JsException(v8ctx, "Replace requires messageId and message arguments")
+		}
+
+		// Get message ID (first argument)
+		if !args[0].IsString() {
+			return bridge.JsException(v8ctx, "messageId must be a string")
+		}
+		messageID := args[0].String()
+
+		// Parse message argument (second argument)
+		msg, err := parseMessage(v8ctx, args[1])
+		if err != nil {
+			return bridge.JsException(v8ctx, "invalid message: "+err.Error())
+		}
+
+		// Set message ID to the provided ID
+		msg.MessageID = messageID
+
+		// Set delta mode for replacement
+		msg.Delta = true
+		msg.DeltaAction = message.DeltaReplace
+		msg.DeltaPath = "" // Empty path means replace entire message
+
+		// Call ctx.Send
+		if err := ctx.Send(msg); err != nil {
+			return bridge.JsException(v8ctx, "Replace failed: "+err.Error())
+		}
+
+		// Automatically flush after sending
+		if err := ctx.Flush(); err != nil {
+			return bridge.JsException(v8ctx, "Flush failed: "+err.Error())
+		}
+
+		// Return the message ID
+		returnID, err := v8go.NewValue(iso, messageID)
+		if err != nil {
+			return bridge.JsException(v8ctx, "Failed to create return value: "+err.Error())
+		}
+		return returnID
 	})
 }