Skip to main content

a3s_code_core/planning/
llm_planner.rs

1//! LLM-powered planning logic
2//!
3//! Provides intelligent plan generation, goal extraction, and achievement
4//! evaluation by sending structured prompts to an LLM and parsing JSON responses.
5//! Falls back to heuristic logic when no LLM client is available.
6
7use crate::llm::structured::{generate_blocking, StructuredMode, StructuredRequest};
8use crate::llm::{LlmClient, Message};
9use crate::planning::{AgentGoal, Complexity, ExecutionPlan, Task};
10use anyhow::{Context, Result};
11use serde::{Deserialize, Serialize};
12use std::sync::Arc;
13
14/// Result of evaluating goal achievement
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct AchievementResult {
17    /// Whether the goal has been achieved
18    pub achieved: bool,
19    /// Progress toward goal (0.0 - 1.0)
20    pub progress: f32,
21    /// Criteria that remain unmet
22    pub remaining_criteria: Vec<String>,
23}
24
25/// Pre-analysis result — intent, goal, plan, and optimized input in one LLM call.
26#[derive(Debug, Clone)]
27pub struct PreAnalysis {
28    pub intent: crate::prompts::AgentStyle,
29    pub requires_planning: bool,
30    pub goal: AgentGoal,
31    pub execution_plan: ExecutionPlan,
32    /// LLM-rewritten version of the user input with ambiguities resolved.
33    pub optimized_input: String,
34}
35
36/// LLM-powered planner that generates plans, extracts goals, and evaluates achievement
37pub struct LlmPlanner;
38
39// ============================================================================
40// JSON response schemas for LLM parsing
41// ============================================================================
42
43#[derive(Debug, Deserialize)]
44struct PlanResponse {
45    goal: String,
46    complexity: String,
47    steps: Vec<StepResponse>,
48    #[serde(default)]
49    required_tools: Vec<String>,
50}
51
52#[derive(Debug, Deserialize)]
53struct StepResponse {
54    id: String,
55    description: String,
56    #[serde(default)]
57    tool: Option<String>,
58    #[serde(default)]
59    dependencies: Vec<String>,
60    #[serde(default)]
61    success_criteria: Option<String>,
62}
63
64#[derive(Debug, Deserialize)]
65struct GoalResponse {
66    description: String,
67    success_criteria: Vec<String>,
68}
69
70#[derive(Debug, Deserialize)]
71struct AchievementResponse {
72    achieved: bool,
73    progress: f32,
74    #[serde(default)]
75    remaining_criteria: Vec<String>,
76}
77
78#[derive(Debug, Deserialize)]
79struct PreAnalysisResponse {
80    intent: String,
81    requires_planning: bool,
82    goal: GoalResponse,
83    execution_plan: PreAnalysisPlan,
84    optimized_input: String,
85}
86
87#[derive(Debug, Deserialize)]
88struct PreAnalysisPlan {
89    complexity: String,
90    steps: Vec<StepResponse>,
91    #[serde(default)]
92    required_tools: Vec<String>,
93}
94
95impl LlmPlanner {
96    /// Generate an execution plan from a prompt using LLM
97    pub async fn create_plan(
98        llm: &Arc<dyn LlmClient>,
99        prompt: &str,
100        language: Option<&str>,
101    ) -> Result<ExecutionPlan> {
102        let system = planning_system_prompt(crate::prompts::LLM_PLAN_SYSTEM, language);
103
104        let messages = vec![Message::user(prompt)];
105        let response = llm
106            .complete(&messages, Some(system.as_str()), &[])
107            .await
108            .context("LLM call failed during plan creation")?;
109
110        let text = response.text();
111        Self::parse_plan_response(&text)
112    }
113
114    /// Extract a goal with success criteria from a prompt using LLM
115    pub async fn extract_goal(
116        llm: &Arc<dyn LlmClient>,
117        prompt: &str,
118        language: Option<&str>,
119    ) -> Result<AgentGoal> {
120        let system = planning_system_prompt(crate::prompts::LLM_GOAL_EXTRACT_SYSTEM, language);
121
122        let messages = vec![Message::user(prompt)];
123        let response = llm
124            .complete(&messages, Some(system.as_str()), &[])
125            .await
126            .context("LLM call failed during goal extraction")?;
127
128        let text = response.text();
129        Self::parse_goal_response(&text)
130    }
131
132    /// Evaluate whether a goal has been achieved given current state
133    pub async fn check_achievement(
134        llm: &Arc<dyn LlmClient>,
135        goal: &AgentGoal,
136        current_state: &str,
137    ) -> Result<AchievementResult> {
138        let prompt = format!(
139            "Goal: {}\nSuccess Criteria: {}\nCurrent State: {}",
140            goal.description,
141            goal.success_criteria.join("; "),
142            current_state,
143        );
144        let req = StructuredRequest {
145            prompt,
146            system: Some(crate::prompts::LLM_GOAL_CHECK_SYSTEM.to_string()),
147            schema: Self::achievement_schema(),
148            schema_name: "goal_achievement".to_string(),
149            schema_description: Some(
150                "Strict, evidence-backed evaluation of whether every goal criterion is met"
151                    .to_string(),
152            ),
153            mode: StructuredMode::Auto,
154            max_repair_attempts: 2,
155        };
156
157        let result = generate_blocking(&**llm, &req)
158            .await
159            .context("LLM achievement structured generation failed")?;
160        Self::achievement_from_value(result.object)
161    }
162
163    /// Create a minimal fallback plan when explicit planning cannot use the LLM.
164    ///
165    /// The original request is the only honest executable step available here.
166    /// Fabricating numbered placeholder steps makes the task tracker look active
167    /// without conveying useful progress.
168    pub fn fallback_plan(prompt: &str) -> ExecutionPlan {
169        let content = match prompt.trim() {
170            "" => "Complete the requested task",
171            prompt => prompt,
172        };
173        let mut plan = ExecutionPlan::new(content, Complexity::Simple);
174        plan.add_step(Task::new("step-1", content));
175
176        plan
177    }
178
179    /// Create a fallback goal using heuristic logic (no LLM required).
180    ///
181    /// Do not invent English success criteria — they surface in product UI and
182    /// fight the host language pin. The request text is the only honest goal.
183    pub fn fallback_goal(prompt: &str) -> AgentGoal {
184        AgentGoal::new(prompt)
185    }
186
187    /// Create a fail-closed fallback achievement result (no LLM required).
188    ///
189    /// Goal completion is a control-plane decision: a host may stop an
190    /// unbounded engineering loop as soon as `GoalAchieved` is emitted. Plain
191    /// answer text such as "done" (or even "not complete") is not sufficient
192    /// evidence for that decision. If the structured evaluator is unavailable,
193    /// keep every success criterion open and let the host retry in a later
194    /// iteration.
195    pub fn fallback_check_achievement(goal: &AgentGoal, current_state: &str) -> AchievementResult {
196        let _ = current_state;
197
198        AchievementResult {
199            achieved: false,
200            progress: goal.progress,
201            remaining_criteria: goal.success_criteria.clone(),
202        }
203    }
204
205    /// Perform pre-analysis in a single LLM call: intent classification, goal extraction,
206    /// execution plan, and input optimization. Falls back to heuristics on failure.
207    pub async fn pre_analyze(
208        llm: &Arc<dyn LlmClient>,
209        prompt: &str,
210        language: Option<&str>,
211    ) -> Result<PreAnalysis> {
212        let system = planning_system_prompt(crate::prompts::PRE_ANALYSIS_SYSTEM, language);
213        let req = StructuredRequest {
214            prompt: format!(
215                "Analyze this user request and return a compact pre-analysis object. \
216                 Use at most 5 execution steps.\n\nUser request:\n{prompt}"
217            ),
218            system: Some(system),
219            schema: Self::pre_analysis_schema(),
220            schema_name: "pre_analysis".to_string(),
221            schema_description: Some(
222                "Intent, goal, plan, and optimized input for an agent turn".to_string(),
223            ),
224            mode: StructuredMode::Auto,
225            max_repair_attempts: 2,
226        };
227
228        let result = generate_blocking(&**llm, &req)
229            .await
230            .context("LLM pre-analysis structured generation failed")?;
231
232        Self::pre_analysis_from_value(result.object, prompt)
233            .context("Failed to parse pre-analysis JSON from LLM response")
234    }
235
236    fn pre_analysis_from_value(
237        value: serde_json::Value,
238        original_prompt: &str,
239    ) -> Result<PreAnalysis> {
240        let parsed: PreAnalysisResponse = serde_json::from_value(value)
241            .context("pre-analysis object did not match the expected response shape")?;
242        Self::pre_analysis_from_response(parsed, original_prompt)
243    }
244
245    fn pre_analysis_from_response(
246        parsed: PreAnalysisResponse,
247        original_prompt: &str,
248    ) -> Result<PreAnalysis> {
249        let intent = match parsed.intent.to_lowercase().as_str() {
250            "plan" => crate::prompts::AgentStyle::Plan,
251            "explore" => crate::prompts::AgentStyle::Explore,
252            "verification" => crate::prompts::AgentStyle::Verification,
253            "codereview" | "code review" => crate::prompts::AgentStyle::CodeReview,
254            _ => crate::prompts::AgentStyle::GeneralPurpose,
255        };
256
257        let goal_description = parsed.goal.description.clone();
258        let goal =
259            AgentGoal::new(goal_description.clone()).with_criteria(parsed.goal.success_criteria);
260
261        let complexity = match parsed.execution_plan.complexity.as_str() {
262            "Simple" => Complexity::Simple,
263            "Medium" => Complexity::Medium,
264            "Complex" => Complexity::Complex,
265            "VeryComplex" => Complexity::VeryComplex,
266            _ => Complexity::Medium,
267        };
268
269        let mut plan = ExecutionPlan::new(goal_description, complexity);
270        for step_resp in parsed.execution_plan.steps {
271            let mut task = Task::new(step_resp.id, step_resp.description);
272            if let Some(tool) = step_resp.tool {
273                task = task.with_tool(tool);
274            }
275            if !step_resp.dependencies.is_empty() {
276                task = task.with_dependencies(step_resp.dependencies);
277            }
278            if let Some(criteria) = step_resp.success_criteria {
279                task = task.with_success_criteria(criteria);
280            }
281            plan.add_step(task);
282        }
283        for tool in parsed.execution_plan.required_tools {
284            plan.add_required_tool(tool);
285        }
286
287        Ok(PreAnalysis {
288            intent,
289            requires_planning: parsed.requires_planning,
290            goal,
291            execution_plan: plan,
292            optimized_input: if parsed.optimized_input.is_empty() {
293                original_prompt.to_string()
294            } else {
295                parsed.optimized_input
296            },
297        })
298    }
299
300    fn pre_analysis_schema() -> serde_json::Value {
301        serde_json::json!({
302            "type": "object",
303            "required": ["intent", "requires_planning", "goal", "execution_plan", "optimized_input"],
304            "properties": {
305                "intent": { "type": "string" },
306                "requires_planning": { "type": "boolean" },
307                "goal": {
308                    "type": "object",
309                    "required": ["description", "success_criteria"],
310                    "properties": {
311                        "description": { "type": "string", "minLength": 1 },
312                        "success_criteria": {
313                            "type": "array",
314                            "items": { "type": "string" }
315                        }
316                    }
317                },
318                "execution_plan": {
319                    "type": "object",
320                    "required": ["complexity", "steps"],
321                    "properties": {
322                        "complexity": { "type": "string" },
323                        "steps": {
324                            "type": "array",
325                            "items": {
326                                "type": "object",
327                                "required": ["id", "description"],
328                                "properties": {
329                                    "id": { "type": "string" },
330                                    "description": { "type": "string" },
331                                    "tool": { "type": "string" },
332                                    "dependencies": {
333                                        "type": "array",
334                                        "items": { "type": "string" }
335                                    },
336                                    "success_criteria": { "type": "string" }
337                                }
338                            }
339                        },
340                        "required_tools": {
341                            "type": "array",
342                            "items": { "type": "string" }
343                        }
344                    }
345                },
346                "optimized_input": { "type": "string" }
347            }
348        })
349    }
350
351    fn achievement_schema() -> serde_json::Value {
352        serde_json::json!({
353            "type": "object",
354            "additionalProperties": false,
355            "required": ["achieved", "progress", "remaining_criteria"],
356            "properties": {
357                "achieved": { "type": "boolean" },
358                "progress": {
359                    "type": "number",
360                    "minimum": 0.0,
361                    "maximum": 1.0
362                },
363                "remaining_criteria": {
364                    "type": "array",
365                    "items": { "type": "string" }
366                }
367            }
368        })
369    }
370
371    fn achievement_from_value(value: serde_json::Value) -> Result<AchievementResult> {
372        let parsed: AchievementResponse = serde_json::from_value(value)
373            .context("achievement object did not match the expected response shape")?;
374        Ok(AchievementResult {
375            achieved: parsed.achieved,
376            progress: parsed.progress.clamp(0.0, 1.0),
377            remaining_criteria: parsed.remaining_criteria,
378        })
379    }
380
381    // ========================================================================
382    // JSON parsing helpers
383    // ========================================================================
384
385    fn parse_plan_response(text: &str) -> Result<ExecutionPlan> {
386        let parsed: PlanResponse = Self::parse_json_lenient(text)
387            .context("Failed to parse plan JSON from LLM response")?;
388
389        let complexity = match parsed.complexity.as_str() {
390            "Simple" => Complexity::Simple,
391            "Medium" => Complexity::Medium,
392            "Complex" => Complexity::Complex,
393            "VeryComplex" => Complexity::VeryComplex,
394            _ => Complexity::Medium,
395        };
396
397        let mut plan = ExecutionPlan::new(parsed.goal, complexity);
398
399        for step_resp in parsed.steps {
400            let mut task = Task::new(step_resp.id, step_resp.description);
401            if let Some(tool) = step_resp.tool {
402                task = task.with_tool(tool);
403            }
404            if !step_resp.dependencies.is_empty() {
405                task = task.with_dependencies(step_resp.dependencies);
406            }
407            if let Some(criteria) = step_resp.success_criteria {
408                task = task.with_success_criteria(criteria);
409            }
410            plan.add_step(task);
411        }
412
413        for tool in parsed.required_tools {
414            plan.add_required_tool(tool);
415        }
416
417        Ok(plan)
418    }
419
420    fn parse_goal_response(text: &str) -> Result<AgentGoal> {
421        let parsed: GoalResponse = Self::parse_json_lenient(text)
422            .context("Failed to parse goal JSON from LLM response")?;
423
424        Ok(AgentGoal::new(parsed.description).with_criteria(parsed.success_criteria))
425    }
426
427    #[cfg(test)]
428    fn parse_achievement_response(text: &str) -> Result<AchievementResult> {
429        let parsed: AchievementResponse = Self::parse_json_lenient(text)
430            .context("Failed to parse achievement JSON from LLM response")?;
431
432        Ok(AchievementResult {
433            achieved: parsed.achieved,
434            progress: parsed.progress.clamp(0.0, 1.0),
435            remaining_criteria: parsed.remaining_criteria,
436        })
437    }
438
439    /// Parse JSON from possibly-dirty LLM output into the target type.
440    ///
441    /// Uses the shared robust extractor ([`crate::llm::structured::extract_json_value`]),
442    /// which handles ```json fences, surrounding prose, and braces embedded in
443    /// strings — unlike the previous naive first-`{`/last-`}` slice, which broke
444    /// on fenced output or any `}` inside a string value.
445    fn parse_json_lenient<T: serde::de::DeserializeOwned>(text: &str) -> Result<T> {
446        let value = crate::llm::structured::extract_json_value(text)?;
447        Ok(serde_json::from_value(value)?)
448    }
449}
450
451fn planning_system_prompt(base: &str, language: Option<&str>) -> String {
452    match language.map(str::trim).filter(|value| !value.is_empty()) {
453        Some(language) => format!(
454            "{}\n\n{}",
455            base.trim_end(),
456            crate::prompts::output_language_contract(language)
457        ),
458        None => base.to_owned(),
459    }
460}
461
462// ============================================================================
463// Tests
464// ============================================================================
465
466#[cfg(test)]
467mod tests {
468    use super::*;
469
470    #[test]
471    fn planning_system_prompt_appends_shared_language_contract() {
472        let prompt = planning_system_prompt("Base planner.", Some("zh-CN"));
473        assert!(prompt.starts_with("Base planner."));
474        assert!(prompt.contains("## Output Language"));
475        assert!(prompt.contains("zh-CN"));
476        assert_eq!(
477            planning_system_prompt("Base planner.", None),
478            "Base planner."
479        );
480    }
481
482    #[test]
483    fn test_parse_plan_response() {
484        let json = r#"{
485            "goal": "Build a REST API",
486            "complexity": "Complex",
487            "steps": [
488                {
489                    "id": "step-1",
490                    "description": "Set up project structure",
491                    "tool": "bash",
492                    "dependencies": [],
493                    "success_criteria": "Project directory created"
494                },
495                {
496                    "id": "step-2",
497                    "description": "Implement endpoints",
498                    "tool": "write",
499                    "dependencies": ["step-1"],
500                    "success_criteria": "Endpoints respond correctly"
501                }
502            ],
503            "required_tools": ["bash", "write", "read"]
504        }"#;
505
506        let plan = LlmPlanner::parse_plan_response(json).unwrap();
507        assert_eq!(plan.goal, "Build a REST API");
508        assert_eq!(plan.complexity, Complexity::Complex);
509        assert_eq!(plan.steps.len(), 2);
510        assert_eq!(plan.steps[0].id, "step-1");
511        assert_eq!(plan.steps[0].tool, Some("bash".to_string()));
512        assert_eq!(plan.steps[1].dependencies, vec!["step-1".to_string()]);
513        assert_eq!(plan.required_tools, vec!["bash", "write", "read"]);
514    }
515
516    #[test]
517    fn test_parse_plan_response_with_markdown_fences() {
518        let json = "```json\n{\"goal\": \"Test\", \"complexity\": \"Simple\", \"steps\": [{\"id\": \"step-1\", \"description\": \"Do it\"}], \"required_tools\": []}\n```";
519
520        let plan = LlmPlanner::parse_plan_response(json).unwrap();
521        assert_eq!(plan.goal, "Test");
522        assert_eq!(plan.complexity, Complexity::Simple);
523        assert_eq!(plan.steps.len(), 1);
524    }
525
526    #[test]
527    fn test_parse_plan_response_invalid() {
528        let bad_json = "This is not JSON at all";
529        let result = LlmPlanner::parse_plan_response(bad_json);
530        assert!(result.is_err());
531    }
532
533    #[test]
534    fn test_parse_plan_response_unknown_complexity() {
535        let json =
536            r#"{"goal": "Test", "complexity": "Unknown", "steps": [], "required_tools": []}"#;
537        let plan = LlmPlanner::parse_plan_response(json).unwrap();
538        assert_eq!(plan.complexity, Complexity::Medium); // falls back to Medium
539    }
540
541    #[test]
542    fn test_parse_goal_response() {
543        let json = r#"{
544            "description": "Deploy the application to production",
545            "success_criteria": [
546                "All tests pass",
547                "Application is accessible at production URL",
548                "Health check returns 200"
549            ]
550        }"#;
551
552        let goal = LlmPlanner::parse_goal_response(json).unwrap();
553        assert_eq!(goal.description, "Deploy the application to production");
554        assert_eq!(goal.success_criteria.len(), 3);
555        assert_eq!(goal.success_criteria[0], "All tests pass");
556    }
557
558    #[test]
559    fn test_parse_goal_response_invalid() {
560        let result = LlmPlanner::parse_goal_response("not json");
561        assert!(result.is_err());
562    }
563
564    #[test]
565    fn test_parse_achievement_response() {
566        let json = r#"{
567            "achieved": false,
568            "progress": 0.65,
569            "remaining_criteria": ["Health check not verified"]
570        }"#;
571
572        let result = LlmPlanner::parse_achievement_response(json).unwrap();
573        assert!(!result.achieved);
574        assert!((result.progress - 0.65).abs() < f32::EPSILON);
575        assert_eq!(result.remaining_criteria, vec!["Health check not verified"]);
576    }
577
578    #[test]
579    fn test_parse_achievement_response_achieved() {
580        let json = r#"{"achieved": true, "progress": 1.0, "remaining_criteria": []}"#;
581        let result = LlmPlanner::parse_achievement_response(json).unwrap();
582        assert!(result.achieved);
583        assert!((result.progress - 1.0).abs() < f32::EPSILON);
584        assert!(result.remaining_criteria.is_empty());
585    }
586
587    #[test]
588    fn test_parse_achievement_response_clamps_progress() {
589        let json = r#"{"achieved": false, "progress": 1.5, "remaining_criteria": []}"#;
590        let result = LlmPlanner::parse_achievement_response(json).unwrap();
591        assert!((result.progress - 1.0).abs() < f32::EPSILON);
592    }
593
594    #[test]
595    fn test_fallback_plan() {
596        let short_prompt = "Fix bug";
597        let plan = LlmPlanner::fallback_plan(short_prompt);
598        assert_eq!(plan.complexity, Complexity::Simple);
599        assert_eq!(plan.steps.len(), 1);
600        assert_eq!(plan.goal, short_prompt);
601        assert_eq!(plan.steps[0].content, short_prompt);
602
603        let long_prompt = "Implement a comprehensive authentication system with OAuth2 support, JWT tokens, refresh token rotation, multi-factor authentication, and role-based access control across all API endpoints with proper audit logging and session management capabilities for both web and mobile clients, including password reset flows, account lockout policies, and integration with external identity providers such as Google, GitHub, and SAML-based enterprise SSO systems";
604        let plan = LlmPlanner::fallback_plan(long_prompt);
605        assert_eq!(plan.complexity, Complexity::Simple);
606        assert_eq!(plan.steps.len(), 1);
607        assert_eq!(plan.steps[0].content, long_prompt);
608        assert!(
609            !plan.steps[0].content.contains("Execute step"),
610            "fallback plans must not expose placeholder task text"
611        );
612
613        let plan = LlmPlanner::fallback_plan("   ");
614        assert_eq!(plan.goal, "Complete the requested task");
615        assert_eq!(plan.steps[0].content, "Complete the requested task");
616    }
617
618    #[test]
619    fn test_fallback_goal() {
620        let goal = LlmPlanner::fallback_goal("Fix the login bug");
621        assert_eq!(goal.description, "Fix the login bug");
622        assert!(
623            goal.success_criteria.is_empty(),
624            "fallback goals must not invent English success criteria"
625        );
626    }
627
628    #[test]
629    fn test_fallback_check_achievement_fails_closed_on_done_text() {
630        let goal = AgentGoal::new("Test task").with_criteria(vec!["Criterion 1".to_string()]);
631
632        let result = LlmPlanner::fallback_check_achievement(&goal, "The task is done.");
633        assert!(!result.achieved);
634        assert!(result.progress.abs() < f32::EPSILON);
635        assert_eq!(result.remaining_criteria, vec!["Criterion 1"]);
636    }
637
638    #[test]
639    fn test_fallback_check_achievement_rejects_negated_completion_text() {
640        let goal = AgentGoal::new("Test task").with_criteria(vec!["Criterion 1".to_string()]);
641
642        for state in [
643            "The task is not complete.",
644            "Verification is not done.",
645            "The implementation is unfinished.",
646        ] {
647            let result = LlmPlanner::fallback_check_achievement(&goal, state);
648            assert!(!result.achieved, "must fail closed for {state:?}");
649            assert_eq!(result.remaining_criteria, vec!["Criterion 1"]);
650        }
651    }
652
653    #[test]
654    fn test_fallback_check_achievement_not_done() {
655        let goal = AgentGoal::new("Test task")
656            .with_criteria(vec!["Criterion 1".to_string(), "Criterion 2".to_string()]);
657
658        let result = LlmPlanner::fallback_check_achievement(&goal, "Work in progress");
659        assert!(!result.achieved);
660        assert_eq!(result.remaining_criteria.len(), 2);
661    }
662
663    #[test]
664    fn test_parse_json_lenient_plain() {
665        let v: serde_json::Value = LlmPlanner::parse_json_lenient("  {\"a\": 1}  ").unwrap();
666        assert_eq!(v["a"], 1);
667    }
668
669    #[test]
670    fn test_parse_json_lenient_with_fences() {
671        let text = "```json\n{\"a\": 1}\n```";
672        let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
673        assert_eq!(v["a"], 1);
674    }
675
676    #[test]
677    fn test_parse_json_lenient_with_surrounding_prose() {
678        let text = "Here is the plan:\n{\"goal\": \"test\"}\nDone.";
679        let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
680        assert_eq!(v["goal"], "test");
681    }
682
683    #[test]
684    fn test_parse_json_lenient_brace_inside_string_value() {
685        // The old naive first-`{`/last-`}` slice broke when a string value
686        // contained a `}` followed by trailing prose; the robust extractor
687        // balances braces while respecting string boundaries.
688        let text = "Result: {\"note\": \"use a closing brace } here\"} -- end.";
689        let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
690        assert_eq!(v["note"], "use a closing brace } here");
691    }
692
693    #[test]
694    fn test_parse_json_lenient_fenced_with_trailing_prose() {
695        // ```json fence followed by an explanation. The naive parser's
696        // `rfind('}')` could grab a brace from the trailing prose.
697        let text = "```json\n{\"goal\": \"ship\"}\n```\nNote: revisit the `plan` later.";
698        let v: serde_json::Value = LlmPlanner::parse_json_lenient(text).unwrap();
699        assert_eq!(v["goal"], "ship");
700    }
701
702    #[test]
703    fn test_parse_json_lenient_rejects_non_json() {
704        let err = LlmPlanner::parse_json_lenient::<serde_json::Value>("no json here at all");
705        assert!(err.is_err());
706    }
707
708    /// Replays a fixed sequence of assistant text responses, one per call.
709    struct ReplayClient {
710        responses: std::sync::Mutex<Vec<String>>,
711    }
712
713    impl ReplayClient {
714        fn new(responses: Vec<String>) -> Self {
715            Self {
716                responses: std::sync::Mutex::new(responses),
717            }
718        }
719    }
720
721    #[async_trait::async_trait]
722    impl LlmClient for ReplayClient {
723        async fn complete(
724            &self,
725            _messages: &[Message],
726            _system: Option<&str>,
727            _tools: &[crate::llm::ToolDefinition],
728        ) -> anyhow::Result<crate::llm::LlmResponse> {
729            let text = {
730                let mut r = self.responses.lock().unwrap();
731                if r.is_empty() {
732                    String::new()
733                } else {
734                    r.remove(0)
735                }
736            };
737            Ok(crate::llm::LlmResponse {
738                message: Message {
739                    role: "assistant".to_string(),
740                    content: vec![crate::llm::ContentBlock::Text { text }],
741                    reasoning_content: None,
742                },
743                usage: crate::llm::TokenUsage::default(),
744                stop_reason: None,
745                token_logprobs: Vec::new(),
746                meta: None,
747            })
748        }
749
750        async fn complete_streaming(
751            &self,
752            _messages: &[Message],
753            _system: Option<&str>,
754            _tools: &[crate::llm::ToolDefinition],
755            _cancel_token: tokio_util::sync::CancellationToken,
756        ) -> anyhow::Result<tokio::sync::mpsc::Receiver<crate::llm::StreamEvent>> {
757            anyhow::bail!("streaming not used in planner tests")
758        }
759    }
760
761    #[tokio::test]
762    async fn test_pre_analyze_repairs_invalid_json() {
763        // First response is unparseable; pre_analyze must re-prompt and succeed
764        // on the second (valid) response.
765        let good = r#"{"intent":"explore","requires_planning":false,"goal":{"description":"Do x","success_criteria":["done"]},"execution_plan":{"complexity":"Simple","steps":[],"required_tools":[]},"optimized_input":"Do x carefully"}"#;
766        let client: Arc<dyn LlmClient> = Arc::new(ReplayClient::new(vec![
767            "Sorry — here's the plan, but not as JSON.".to_string(),
768            good.to_string(),
769        ]));
770        let pa = LlmPlanner::pre_analyze(&client, "do x", None)
771            .await
772            .unwrap();
773        assert_eq!(pa.optimized_input, "Do x carefully");
774    }
775
776    #[tokio::test]
777    async fn test_pre_analyze_first_try_with_fenced_json() {
778        // A single ```json-fenced response must parse on the first attempt
779        // (robust extractor), with no repair round needed.
780        let good = format!(
781            "```json\n{}\n```",
782            r#"{"intent":"plan","requires_planning":true,"goal":{"description":"g","success_criteria":[]},"execution_plan":{"complexity":"Medium","steps":[],"required_tools":[]},"optimized_input":"opt"}"#
783        );
784        let client: Arc<dyn LlmClient> = Arc::new(ReplayClient::new(vec![good]));
785        let pa = LlmPlanner::pre_analyze(&client, "do x", None)
786            .await
787            .unwrap();
788        assert_eq!(pa.optimized_input, "opt");
789    }
790}