Merge pull request #1546 from trheyi/main

Enhance message handling with execute family and question types

Author: trheyi

agent/assistant/handlers/stream.go

@@ -294,14 +294,24 @@ func (s *streamState) handleToolCall(data []byte) int {
 	return 0 // Continue
 }
 
+// isExecuteFamily checks if a message type belongs to the execute family
+// (execute itself or any semantic sub-type: agent, todo, plan, job).
+// These all share the same two-phase lifecycle (running → completed/error).
+func isExecuteFamily(t string) bool {
+	switch t {
+	case message.TypeExecute, message.TypeAgent, message.TypeTodo, message.TypePlan, message.TypeJob, message.TypeQuestion:
+		return true
+	}
+	return false
+}
+
 // handleExecute handles execute observation chunks from sandbox CLI agents.
 // These represent tool actions observed inside the agent runtime (e.g., Bash, Read, Write).
 func (s *streamState) handleExecute(data []byte) int {
 	if len(data) == 0 {
 		return 0
 	}
 
-	s.currentType = message.TypeExecute
 	s.buffer = append(s.buffer, data...)
 	s.chunkCount++
 	s.messageSeq++
@@ -322,14 +332,18 @@ func (s *streamState) handleExecute(data []byte) int {
 		s.lastExecProps[k] = v
 	}
 
-	deltaAction := "merge"
+	msgType := message.TypeExecute
+	if st, ok := props["semantic_type"].(string); ok && st != "" {
+		msgType = st
+	}
+	s.currentType = msgType
 
 	msg := &message.Message{
 		ChunkID:     s.ctx.IDGenerator.GenerateChunkID(),
 		MessageID:   s.currentGroupID,
-		Type:        message.TypeExecute,
+		Type:        msgType,
 		Delta:       true,
-		DeltaAction: deltaAction,
+		DeltaAction: "merge",
 		Props:       props,
 	}
 
@@ -408,13 +422,14 @@ func (s *streamState) handleMessageEnd(data []byte) int {
 	shouldSkipHistory := s.ctx.Stack != nil && s.ctx.Stack.Options != nil &&
 		s.ctx.Stack.Options.Skip != nil && s.ctx.Stack.Options.Skip.History
 
-	// Execute messages have two (or more) phases sharing the same message_id:
+	// Execute-family messages (execute, agent, todo, plan, job) have two-phase lifecycles:
 	//   1. running / suspended / resumed — streamed for UI display only, NOT persisted
 	//   2. completed / error — the final state, persisted to the buffer
 	// Only persist when we have an explicit terminal status.
-	isExecuteFinal := msgType == message.TypeExecute &&
+	isExecFamily := isExecuteFamily(msgType)
+	isExecuteFinal := isExecFamily &&
 		(s.lastExecStatus == "completed" || s.lastExecStatus == "error")
-	skipExecute := msgType == message.TypeExecute && !isExecuteFinal
+	skipExecute := isExecFamily && !isExecuteFinal
 
 	if s.ctx.Buffer != nil && len(s.buffer) > 0 && !shouldSkipHistory && !skipExecute {
 		assistantID := ""
@@ -423,8 +438,8 @@ func (s *streamState) handleMessageEnd(data []byte) int {
 		}
 
 		var props map[string]interface{}
-		switch msgType {
-		case message.TypeToolCall:
+		switch {
+		case msgType == message.TypeToolCall:
 			var toolCallData interface{}
 			if err := jsoniter.Unmarshal(s.buffer, &toolCallData); err == nil {
 				props = map[string]interface{}{
@@ -435,16 +450,10 @@ func (s *streamState) handleMessageEnd(data []byte) int {
 					"content": string(s.buffer),
 				}
 			}
-		case message.TypeExecute:
-			if s.lastExecProps != nil {
-				props = make(map[string]interface{}, len(s.lastExecProps))
-				for k, v := range s.lastExecProps {
-					props[k] = v
-				}
-			} else {
-				props = map[string]interface{}{
-					"content": string(s.buffer),
-				}
+		case s.lastExecProps != nil:
+			props = make(map[string]interface{}, len(s.lastExecProps))
+			for k, v := range s.lastExecProps {
+				props[k] = v
 			}
 		default:
 			props = map[string]interface{}{

agent/output/builtin.go

@@ -156,7 +156,7 @@ func NewVideoMessage(url string) *message.Message {
 // IsBuiltinType checks if a message type is a built-in type
 func IsBuiltinType(msgType string) bool {
 	switch msgType {
-	case message.TypeUserInput, message.TypeText, message.TypeThinking, message.TypeLoading, message.TypeToolCall, message.TypeExecute, message.TypeError, message.TypeImage, message.TypeAudio, message.TypeVideo, message.TypeAction, message.TypeEvent:
+	case message.TypeUserInput, message.TypeText, message.TypeThinking, message.TypeLoading, message.TypeToolCall, message.TypeExecute, message.TypeAgent, message.TypeTodo, message.TypePlan, message.TypeJob, message.TypeQuestion, message.TypeError, message.TypeImage, message.TypeAudio, message.TypeVideo, message.TypeAction, message.TypeEvent:
 		return true
 	default:
 		return false

agent/output/message/types.go

@@ -78,6 +78,13 @@ const (
 	// Agent execution observation types
 	TypeExecute = "execute" // Agent tool execution observation (sandbox CLI agent actions, not LLM tool_call requests)
 
+	// Semantic execute sub-types (runner-annotated, rendered by dedicated frontend components)
+	TypeAgent    = "agent"    // Sub-agent dispatch (Claude Agent tool, A2A calls)
+	TypeTodo     = "todo"     // Structured task checklist (Claude TodoWrite tool)
+	TypePlan     = "plan"     // Plan mode transitions (Claude EnterPlanMode/ExitPlanMode)
+	TypeJob      = "job"      // Async background jobs (Claude TaskCreate/TaskGet/... — Unix job semantics)
+	TypeQuestion = "question" // Interactive user questions (Claude AskUserQuestion)
+
 	// System types (not visible in standard chat clients)
 	TypeAction = "action" // System action (open panel, navigate, etc.) - silent in OpenAI clients
 	TypeEvent  = "event"  // Lifecycle event (stream_start, stream_end, etc.) - CUI only, silent in OpenAI clients

agent/sandbox/v2/claude/parse.go

@@ -376,12 +376,14 @@ func (p *streamParser) onContentBlockStart(event map[string]any) (stopped bool)
 		return false
 	}
 
-	return p.emitExecute(map[string]any{
+	execProps := map[string]any{
 		"tool":    toolName,
 		"tool_id": toolID,
 		"status":  "running",
 		"runner":  "claude-cli",
-	})
+	}
+	injectSemanticType(execProps, toolName)
+	return p.emitExecute(execProps)
 }
 
 func (p *streamParser) onContentBlockStop() (stopped bool) {
@@ -489,14 +491,16 @@ func (p *streamParser) handleAssistant(msg map[string]any) (stopped bool) {
 			summary := extractSummary(toolName, inputStr)
 			p.toolSummaries[toolID] = summary
 
-			if p.emitExecute(map[string]any{
+			execProps := map[string]any{
 				"tool":    toolName,
 				"tool_id": toolID,
 				"input":   json.RawMessage(inputRaw),
 				"summary": summary,
 				"status":  "running",
 				"runner":  "claude-cli",
-			}) {
+			}
+			injectSemanticType(execProps, toolName)
+			if p.emitExecute(execProps) {
 				p.endMessage()
 				return true
 			}
@@ -572,15 +576,18 @@ func (p *streamParser) handleUser(msg map[string]any) (stopped bool) {
 			"status":   status,
 			"is_error": isError,
 		}
+		toolName := ""
 		if name, ok := p.toolNames[toolUseID]; ok {
 			execProps["tool"] = name
+			toolName = name
 		}
 		if input, ok := p.toolInputs[toolUseID]; ok {
 			execProps["input"] = json.RawMessage(input)
 		}
 		if summary, ok := p.toolSummaries[toolUseID]; ok {
 			execProps["summary"] = summary
 		}
+		injectSemanticType(execProps, toolName)
 
 		if reuseMsgID, ok := p.toolMsgIDs[toolUseID]; ok {
 			if p.beginMessageWithID(reuseMsgID, "execute") {
@@ -653,3 +660,44 @@ func (p *streamParser) handleError(msg map[string]any) error {
 	}
 	return nil
 }
+
+// resolveSemanticType maps Claude Code CLI tool names to semantic types.
+// Returns ("", nil) for tools that should remain as "execute".
+func resolveSemanticType(tool string) (string, map[string]any) {
+	switch tool {
+	case "Agent":
+		return message.TypeAgent, nil
+	case "TodoWrite":
+		return message.TypeTodo, map[string]any{"action": "write"}
+	case "EnterPlanMode":
+		return message.TypePlan, map[string]any{"action": "enter"}
+	case "ExitPlanMode":
+		return message.TypePlan, map[string]any{"action": "exit"}
+	case "TaskCreate":
+		return message.TypeJob, map[string]any{"action": "create"}
+	case "TaskGet":
+		return message.TypeJob, map[string]any{"action": "get"}
+	case "TaskList":
+		return message.TypeJob, map[string]any{"action": "list"}
+	case "TaskOutput":
+		return message.TypeJob, map[string]any{"action": "output"}
+	case "TaskStop":
+		return message.TypeJob, map[string]any{"action": "stop"}
+	case "TaskUpdate":
+		return message.TypeJob, map[string]any{"action": "update"}
+	case "AskUserQuestion":
+		return message.TypeQuestion, nil
+	}
+	return "", nil
+}
+
+// injectSemanticType annotates execute props with semantic_type and extra metadata.
+func injectSemanticType(props map[string]any, toolName string) {
+	semanticType, extraProps := resolveSemanticType(toolName)
+	if semanticType != "" {
+		props["semantic_type"] = semanticType
+		for k, v := range extraProps {
+			props[k] = v
+		}
+	}
+}

agent/sandbox/v2/claude/runner.go

@@ -61,7 +61,7 @@ func (r *Runner) Prepare(ctx context.Context, req *types.PrepareRequest) error {
 		if err := shared.InjectSystemSkills(ws, tools.SkillsFS, ".claude/skills"); err != nil {
 			r.logger.Warn("inject system skills: %v", err)
 		}
-		if err := shared.InjectAgentDefinitions(ws, tools.AgentsFS, ".claude/agents"); err != nil {
+		if err := shared.InjectAgentDefinitions(ws, tools.AgentsFS, prefix+"/agents"); err != nil {
 			r.logger.Warn("inject agent definitions: %v", err)
 		}
 		if err := shared.AppendSystemPrompt(ws, "CLAUDE.md", tools.SystemPrompt); err != nil {

tools/agent/call_schema.json

@@ -1,6 +1,7 @@
 {
   "name": "agent_call",
   "description": "Call another AI expert agent with a message and get the response. Use when the user @mentions an expert or when you need specialized help from another agent.",
+  "process": "tools.agent_call",
   "inputSchema": {
     "type": "object",
     "properties": {

tools/prompts/system-tools.md

@@ -110,7 +110,7 @@ The system skills (`yao-web`, `yao-process`, `yao-doc`, `yao-image`, `yao-agent`
 
 User messages may contain `<Mention>` tags referencing experts, workspaces, files, or directories:
 
-- `<Mention type="expert" value="assistant_id">Name</Mention>` — The user wants to involve this AI expert. Use the Agent tool with `subagent_type="a2a"` to delegate the task, passing the `assistant_id` and relevant context in the prompt.
+- `<Mention type="expert" value="assistant_id">Name</Mention>` — The user wants to involve this AI expert. Use `tai tool agent_call '{"assistant_id":"<the value>","message":"<relevant query>"}'` to call the expert and get their response.
 - `<Mention type="workspace" value="workspace_id">Name</Mention>` — References a workspace. Use `workspace_file_list` and `workspace_file_read` to access its files.
 - `<Mention type="file" value="workspace://wsId/path">Filename</Mention>` — References a specific file. Use `workspace_file_read` to read its content.
 - `<Mention type="directory" value="workspace://wsId/path">DirName</Mention>` — References a directory. Use `workspace_file_list` to browse its contents first, then `workspace_file_read` for specific files.