Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions components/backend/websocket/agui_proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -1128,17 +1128,19 @@ func updateLastActivityTime(projectName, sessionName string, immediate bool) {
}()
}

// isAskUserQuestionToolCall checks if a tool call name is the AskUserQuestion HITL tool.
// Uses case-insensitive comparison after stripping non-alpha characters,
// matching the frontend pattern in use-agent-status.ts.
func isAskUserQuestionToolCall(name string) bool {
// isHITLToolCall checks if a tool call name is a HITL (human-in-the-loop) tool.
// Matches AskUserQuestion and ExitPlanMode using case-insensitive comparison
// after stripping non-alpha characters, matching the frontend pattern in
// use-agent-status.ts.
func isHITLToolCall(name string) bool {
var clean strings.Builder
for _, r := range strings.ToLower(name) {
if r >= 'a' && r <= 'z' {
clean.WriteRune(r)
}
}
return clean.String() == "askuserquestion"
normalized := clean.String()
return normalized == "askuserquestion" || normalized == "exitplanmode"
}

// ─── Between-Run Listener ────────────────────────────────────────────
Expand Down
6 changes: 3 additions & 3 deletions components/backend/websocket/agui_store.go
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,7 @@ func DeriveAgentStatus(sessionID string) string {
return types.AgentStatusIdle
}
}
if toolName, _ := evt["toolCallName"].(string); isAskUserQuestionToolCall(toolName) {
if toolName, _ := evt["toolCallName"].(string); isHITLToolCall(toolName) {
return types.AgentStatusWaitingInput
}
}
Expand Down Expand Up @@ -527,9 +527,9 @@ func compactFinishedRun(sessionID string) {
types.EventTypeStepStarted, types.EventTypeStepFinished:
snapshots = append(snapshots, evt)
case types.EventTypeToolCallStart:
// Preserve AskUserQuestion tool calls — DeriveAgentStatus() needs them
// Preserve HITL tool calls — DeriveAgentStatus() needs them
// to detect waiting_input status after compaction.
if toolName, _ := evt["toolCallName"].(string); isAskUserQuestionToolCall(toolName) {
if toolName, _ := evt["toolCallName"].(string); isHITLToolCall(toolName) {
snapshots = append(snapshots, evt)
}
case types.EventTypeRaw, types.EventTypeCustom, types.EventTypeMeta:
Expand Down
37 changes: 34 additions & 3 deletions components/backend/websocket/agui_store_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -186,15 +186,46 @@ func TestDeriveAgentStatus(t *testing.T) {
}
})

t.Run("case-insensitive AskUserQuestion detection", func(t *testing.T) {
t.Run("RUN_FINISHED with same-run ExitPlanMode returns waiting_input", func(t *testing.T) {
sessionID := "test-session-exit-plan-mode"
sessionsDir := filepath.Join(tmpDir, "sessions", sessionID)
if err := os.MkdirAll(sessionsDir, 0755); err != nil {
t.Fatalf("Failed to create sessions dir: %v", err)
}

events := []map[string]interface{}{
{"type": types.EventTypeRunStarted, "runId": "run-123"},
{"type": types.EventTypeToolCallStart, "runId": "run-123", "toolCallName": "ExitPlanMode"},
{"type": types.EventTypeRunFinished, "runId": "run-123"},
}
eventsFile := filepath.Join(sessionsDir, "agui-events.jsonl")
f, err := os.Create(eventsFile)
if err != nil {
t.Fatalf("Failed to create events file: %v", err)
}
for _, evt := range events {
data, _ := json.Marshal(evt)
f.Write(append(data, '\n'))
}
f.Close()

status := DeriveAgentStatus(sessionID)
if status != types.AgentStatusWaitingInput {
t.Errorf("Expected %q for same-run ExitPlanMode, got %q", types.AgentStatusWaitingInput, status)
}
})

t.Run("case-insensitive HITL tool detection", func(t *testing.T) {
sessionID := "test-session-case-insensitive"
sessionsDir := filepath.Join(tmpDir, "sessions", sessionID)
if err := os.MkdirAll(sessionsDir, 0755); err != nil {
t.Fatalf("Failed to create sessions dir: %v", err)
}

// Test various casings of AskUserQuestion
testCases := []string{"askuserquestion", "ASKUSERQUESTION", "AskUserQuestion", "AsKuSeRqUeStIoN"}
testCases := []string{
"askuserquestion", "ASKUSERQUESTION", "AskUserQuestion", "AsKuSeRqUeStIoN",
"exitplanmode", "EXITPLANMODE", "ExitPlanMode", "eXiTpLaNmOdE",
}
for _, toolName := range testCases {
events := []map[string]interface{}{
{"type": types.EventTypeRunStarted, "runId": "run-123"},
Expand Down
207 changes: 207 additions & 0 deletions components/frontend/src/components/session/exit-plan-mode.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,207 @@
"use client";

import React, { useState } from "react";
import ReactMarkdown from "react-markdown";
import remarkGfm from "remark-gfm";
import { cn } from "@/lib/utils";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { ClipboardCheck, CheckCircle2, Send } from "lucide-react";
import { formatTimestamp } from "@/lib/format-timestamp";
import type { ToolUseBlock, ToolResultBlock } from "@/types/agentic-session";

export type ExitPlanModeMessageProps = {
toolUseBlock: ToolUseBlock;
resultBlock?: ToolResultBlock;
timestamp?: string;
onSubmitAnswer?: (formattedAnswer: string) => Promise<void>;
isNewest?: boolean;
};

type AllowedPrompt = {
tool: string;
prompt: string;
};

function hasResult(resultBlock?: ToolResultBlock): boolean {
if (!resultBlock) return false;
const content = resultBlock.content;
if (!content) return false;
if (typeof content === "string" && content.trim() === "") return false;
return true;
}

export const ExitPlanModeMessage: React.FC<ExitPlanModeMessageProps> = ({
toolUseBlock,
resultBlock,
timestamp,
onSubmitAnswer,
isNewest = false,
}) => {
const input = toolUseBlock.input;
const planContent = (input.planContent as string) || "";
const allowedPrompts = (input.allowedPrompts as AllowedPrompt[]) || [];
Comment on lines +35 to +36
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Validate tool payload shapes before rendering

Lines 42-43 assume planContent is string and allowedPrompts is array. If payload shape is malformed, this can crash at Line 117 (map on non-array) or pass invalid input into ReactMarkdown at Line 106.

Proposed fix
-  const planContent = (input.planContent as string) || "";
-  const allowedPrompts = (input.allowedPrompts as AllowedPrompt[]) || [];
+  const planContent =
+    typeof input.planContent === "string" ? input.planContent : "";
+  const allowedPrompts: AllowedPrompt[] = Array.isArray(input.allowedPrompts)
+    ? input.allowedPrompts.filter(
+        (p): p is AllowedPrompt =>
+          !!p &&
+          typeof p === "object" &&
+          typeof (p as { tool?: unknown }).tool === "string" &&
+          typeof (p as { prompt?: unknown }).prompt === "string",
+      )
+    : [];

Also applies to: 104-109, 117-123

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/frontend/src/components/session/exit-plan-mode.tsx` around lines
42 - 43, Validate and coerce the incoming payload before rendering: ensure
planContent is a string (e.g., check typeof input.planContent === "string" and
fallback to "" for ReactMarkdown) and ensure allowedPrompts is an array (use
Array.isArray(input.allowedPrompts) ? input.allowedPrompts as AllowedPrompt[] :
[]), then use that sanitized values for rendering and mapping; update usages in
this component (references: planContent, allowedPrompts, ReactMarkdown and the
map over allowedPrompts) so malformed payloads don’t cause runtime errors.

const alreadyAnswered = hasResult(resultBlock);
const formattedTime = formatTimestamp(timestamp);

const [submitted, setSubmitted] = useState(false);
const [isSubmitting, setIsSubmitting] = useState(false);
const [showFeedback, setShowFeedback] = useState(false);
const [feedback, setFeedback] = useState("");
const disabled = alreadyAnswered || submitted || isSubmitting || !isNewest;

const handleDecision = async (decision: "approve" | "reject" | "request_changes", feedbackText?: string) => {
if (!onSubmitAnswer || disabled) return;
const response: Record<string, string> = { decision };
if (feedbackText) {
response.feedback = feedbackText;
}
try {
setIsSubmitting(true);
await onSubmitAnswer(JSON.stringify(response));
setSubmitted(true);
} finally {
setIsSubmitting(false);
}
};

return (
<div className="mb-3">
<div className="flex items-start gap-3">
{/* Avatar */}
<div className="flex-shrink-0">
<div
className={cn(
"w-8 h-8 rounded-full flex items-center justify-center",
disabled ? "bg-green-600" : "bg-blue-500"
)}
>
{disabled ? (
<CheckCircle2 className="w-4 h-4 text-white" />
) : (
<ClipboardCheck className="w-4 h-4 text-white" />
)}
</div>
</div>

{/* Content */}
<div className="flex-1 min-w-0">
{formattedTime && (
<div className="text-[10px] text-muted-foreground/60 mb-0.5">{formattedTime}</div>
)}

<div
className={cn(
"rounded-lg border-l-3 pl-3 pr-3 py-2.5",
disabled
? "border-l-green-500 bg-green-50/30 dark:bg-green-950/10"
: "border-l-blue-500 bg-blue-50/30 dark:bg-blue-950/10"
)}
>
<p className="text-sm font-medium text-foreground mb-2">Plan Review</p>

{/* Plan content */}
{planContent && (
<div className="text-sm prose prose-sm dark:prose-invert max-w-none mb-3 max-h-96 overflow-y-auto rounded border border-border/40 p-2.5 bg-background/50">
<ReactMarkdown remarkPlugins={[remarkGfm]}>
{planContent}
</ReactMarkdown>
</div>
)}

{/* Allowed prompts */}
{allowedPrompts.length > 0 && (
<div className="mb-3">
<p className="text-xs font-medium text-muted-foreground mb-1">Requested permissions:</p>
<ul className="space-y-0.5">
{allowedPrompts.map((p, i) => (
<li key={i} className="text-xs text-muted-foreground flex items-center gap-1.5">
<span className="inline-block w-1 h-1 rounded-full bg-muted-foreground/50 flex-shrink-0" />
<span className="font-mono">{p.tool}</span>: {p.prompt}
</li>
))}
</ul>
</div>
)}

{/* Request changes feedback input */}
{showFeedback && !disabled && (
<div className="mb-2">
<Input
autoFocus
placeholder="Describe the changes you'd like..."
value={feedback}
onChange={(e) => setFeedback(e.target.value)}
onKeyDown={(e) => {
if (e.key === "Enter" && !e.shiftKey && feedback.trim()) {
e.preventDefault();
handleDecision("request_changes", feedback.trim());
}
}}
disabled={disabled}
className="h-8 text-sm"
/>
</div>
)}

{/* Action buttons */}
{!disabled && (
<div className="flex items-center gap-1.5 mt-2 pt-1.5 border-t border-border/40">
{!showFeedback ? (
<>
<Button
size="sm"
className="h-7 text-xs gap-1 px-3"
onClick={() => handleDecision("approve")}
>
<CheckCircle2 className="w-3 h-3" />
Approve
</Button>
<Button
variant="outline"
size="sm"
className="h-7 text-xs px-3"
onClick={() => handleDecision("reject")}
>
Reject
</Button>
<Button
variant="ghost"
size="sm"
className="h-7 text-xs px-3"
onClick={() => setShowFeedback(true)}
>
Request Changes
</Button>
</>
) : (
<>
<Button
size="sm"
className="h-7 text-xs gap-1 px-3"
onClick={() => handleDecision("request_changes", feedback.trim())}
disabled={!feedback.trim()}
>
<Send className="w-3 h-3" />
Send Feedback
</Button>
<Button
variant="ghost"
size="sm"
className="h-7 text-xs px-3"
onClick={() => setShowFeedback(false)}
>
Cancel
</Button>
</>
)}
</div>
)}
</div>
</div>
</div>
</div>
);
};

ExitPlanModeMessage.displayName = "ExitPlanModeMessage";
19 changes: 19 additions & 0 deletions components/frontend/src/components/ui/stream-message.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { MessageObject, ToolUseMessages, HierarchicalToolMessage } from "@/types
import { LoadingDots, Message } from "@/components/ui/message";
import { ToolMessage } from "@/components/ui/tool-message";
import { AskUserQuestionMessage } from "@/components/session/ask-user-question";
import { ExitPlanModeMessage } from "@/components/session/exit-plan-mode";
import { ThinkingMessage } from "@/components/ui/thinking-message";
import { SystemMessage } from "@/components/ui/system-message";
import { Button } from "@/components/ui/button";
Expand All @@ -25,6 +26,11 @@ function isAskUserQuestionTool(name: string): boolean {
return normalized === "askuserquestion";
}

function isExitPlanModeTool(name: string): boolean {
const normalized = name.toLowerCase().replace(/[^a-z]/g, "");
return normalized === "exitplanmode";
}

const getRandomAgentMessage = () => {
const messages = [
"The agents are working together on your request...",
Expand Down Expand Up @@ -59,6 +65,19 @@ export const StreamMessage: React.FC<StreamMessageProps> = ({ message, onGoToRes
);
}

// Render ExitPlanMode with plan approval component
if (isExitPlanModeTool(message.toolUseBlock.name)) {
return (
<ExitPlanModeMessage
toolUseBlock={message.toolUseBlock}
resultBlock={message.resultBlock}
timestamp={message.timestamp}
onSubmitAnswer={onSubmitAnswer}
isNewest={isNewest}
/>
);
}

// Check if this is a hierarchical message with children
const hierarchical = message as HierarchicalToolMessage;
return (
Expand Down
8 changes: 4 additions & 4 deletions components/frontend/src/hooks/use-agent-status.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ import type {
} from "@/types/agentic-session";
import type { PlatformMessage } from "@/types/agui";

function isAskUserQuestionTool(name: string): boolean {
function isHITLTool(name: string): boolean {
const normalized = name.toLowerCase().replace(/[^a-z]/g, "");
return normalized === "askuserquestion";
return normalized === "askuserquestion" || normalized === "exitplanmode";
}

/**
Expand All @@ -30,15 +30,15 @@ export function useAgentStatus(
// Non-running phases
if (phase !== "Running") return "idle";

// Scan backwards for the last tool call to check for unanswered AskUserQuestion.
// Scan backwards for the last tool call to check for unanswered HITL tools.
// Raw AG-UI messages store tool calls in msg.toolCalls[] (PlatformToolCall[]).
for (let i = messages.length - 1; i >= 0; i--) {
const msg = messages[i];
if (!msg.toolCalls || msg.toolCalls.length === 0) continue;

// Check the last tool call on this message
const lastTc = msg.toolCalls[msg.toolCalls.length - 1];
if (lastTc.function?.name && isAskUserQuestionTool(lastTc.function.name)) {
if (lastTc.function?.name && isHITLTool(lastTc.function.name)) {
const hasResult =
lastTc.result !== undefined &&
lastTc.result !== null &&
Expand Down
Loading
Loading