Zen MCP Server

""" Data models for tool responses and interactions """ from enum import Enum from typing import Any, Literal, Optional from pydantic import BaseModel, Field class ToolModelCategory(Enum): """Categories for tool model selection based on requirements.""" EXTENDED_REASONING = "extended_reasoning" # Requires deep thinking capabilities FAST_RESPONSE = "fast_response" # Speed and cost efficiency preferred BALANCED = "balanced" # Balance of capability and performance class ContinuationOffer(BaseModel): """Offer for CLI agent to continue conversation when Gemini doesn't ask follow-up""" continuation_id: str = Field( ..., description="Thread continuation ID for multi-turn conversations across different tools" ) note: str = Field(..., description="Message explaining continuation opportunity to CLI agent") remaining_turns: int = Field(..., description="Number of conversation turns remaining") class ToolOutput(BaseModel): """Standardized output format for all tools""" status: Literal[ "success", "error", "files_required_to_continue", "full_codereview_required", "focused_review_required", "test_sample_needed", "more_tests_required", "refactor_analysis_complete", "trace_complete", "resend_prompt", "code_too_large", "continuation_available", "no_bug_found", ] = "success" content: Optional[str] = Field(None, description="The main content/response from the tool") content_type: Literal["text", "markdown", "json"] = "text" metadata: Optional[dict[str, Any]] = Field(default_factory=dict) continuation_offer: Optional[ContinuationOffer] = Field( None, description="Optional offer for Agent to continue conversation" ) class FilesNeededRequest(BaseModel): """Request for missing files / code to continue""" status: Literal["files_required_to_continue"] = "files_required_to_continue" mandatory_instructions: str = Field(..., description="Critical instructions for Agent regarding required context") files_needed: Optional[list[str]] = Field( default_factory=list, description="Specific files that are needed for analysis" ) suggested_next_action: Optional[dict[str, Any]] = Field( None, description="Suggested tool call with parameters after getting clarification", ) class FullCodereviewRequired(BaseModel): """Request for full code review when scope is too large for quick review""" status: Literal["full_codereview_required"] = "full_codereview_required" important: Optional[str] = Field(None, description="Important message about escalation") reason: Optional[str] = Field(None, description="Reason why full review is needed") class FocusedReviewRequired(BaseModel): """Request for Agent to provide smaller, focused subsets of code for review""" status: Literal["focused_review_required"] = "focused_review_required" reason: str = Field(..., description="Why the current scope is too large for effective review") suggestion: str = Field( ..., description="Suggested approach for breaking down the review into smaller, focused parts" ) class TestSampleNeeded(BaseModel): """Request for additional test samples to determine testing framework""" status: Literal["test_sample_needed"] = "test_sample_needed" reason: str = Field(..., description="Reason why additional test samples are required") class MoreTestsRequired(BaseModel): """Request for continuation to generate additional tests""" status: Literal["more_tests_required"] = "more_tests_required" pending_tests: str = Field(..., description="List of pending tests to be generated") class RefactorOpportunity(BaseModel): """A single refactoring opportunity with precise targeting information""" id: str = Field(..., description="Unique identifier for this refactoring opportunity") type: Literal["decompose", "codesmells", "modernize", "organization"] = Field( ..., description="Type of refactoring" ) severity: Literal["critical", "high", "medium", "low"] = Field(..., description="Severity level") file: str = Field(..., description="Absolute path to the file") start_line: int = Field(..., description="Starting line number") end_line: int = Field(..., description="Ending line number") context_start_text: str = Field(..., description="Exact text from start line for verification") context_end_text: str = Field(..., description="Exact text from end line for verification") issue: str = Field(..., description="Clear description of what needs refactoring") suggestion: str = Field(..., description="Specific refactoring action to take") rationale: str = Field(..., description="Why this improves the code") code_to_replace: str = Field(..., description="Original code that should be changed") replacement_code_snippet: str = Field(..., description="Refactored version of the code") new_code_snippets: Optional[list[dict]] = Field( default_factory=list, description="Additional code snippets to be added" ) class RefactorAction(BaseModel): """Next action for Agent to implement refactoring""" action_type: Literal["EXTRACT_METHOD", "SPLIT_CLASS", "MODERNIZE_SYNTAX", "REORGANIZE_CODE", "DECOMPOSE_FILE"] = ( Field(..., description="Type of action to perform") ) target_file: str = Field(..., description="Absolute path to target file") source_lines: str = Field(..., description="Line range (e.g., '45-67')") description: str = Field(..., description="Step-by-step action description for CLI Agent") class RefactorAnalysisComplete(BaseModel): """Complete refactor analysis with prioritized opportunities""" status: Literal["refactor_analysis_complete"] = "refactor_analysis_complete" refactor_opportunities: list[RefactorOpportunity] = Field(..., description="List of refactoring opportunities") priority_sequence: list[str] = Field(..., description="Recommended order of refactoring IDs") next_actions: list[RefactorAction] = Field(..., description="Specific actions for the agent to implement") class CodeTooLargeRequest(BaseModel): """Request to reduce file selection due to size constraints""" status: Literal["code_too_large"] = "code_too_large" content: str = Field(..., description="Message explaining the size constraint") content_type: Literal["text"] = "text" metadata: dict[str, Any] = Field(default_factory=dict) class ResendPromptRequest(BaseModel): """Request to resend prompt via file due to size limits""" status: Literal["resend_prompt"] = "resend_prompt" content: str = Field(..., description="Instructions for handling large prompt") content_type: Literal["text"] = "text" metadata: dict[str, Any] = Field(default_factory=dict) class TraceEntryPoint(BaseModel): """Entry point information for trace analysis""" file: str = Field(..., description="Absolute path to the file") class_or_struct: str = Field(..., description="Class or module name") method: str = Field(..., description="Method or function name") signature: str = Field(..., description="Full method signature") parameters: Optional[dict[str, Any]] = Field(default_factory=dict, description="Parameter values used in analysis") class TraceTarget(BaseModel): """Target information for dependency analysis""" file: str = Field(..., description="Absolute path to the file") class_or_struct: str = Field(..., description="Class or module name") method: str = Field(..., description="Method or function name") signature: str = Field(..., description="Full method signature") class CallPathStep(BaseModel): """A single step in the call path trace""" from_info: dict[str, Any] = Field(..., description="Source location information", alias="from") to: dict[str, Any] = Field(..., description="Target location information") reason: str = Field(..., description="Reason for the call or dependency") condition: Optional[str] = Field(None, description="Conditional logic if applicable") ambiguous: bool = Field(False, description="Whether this call is ambiguous") class BranchingPoint(BaseModel): """A branching point in the execution flow""" file: str = Field(..., description="File containing the branching point") method: str = Field(..., description="Method containing the branching point") line: int = Field(..., description="Line number of the branching point") condition: str = Field(..., description="Branching condition") branches: list[str] = Field(..., description="Possible execution branches") ambiguous: bool = Field(False, description="Whether the branching is ambiguous") class SideEffect(BaseModel): """A side effect detected in the trace""" type: str = Field(..., description="Type of side effect") description: str = Field(..., description="Description of the side effect") file: str = Field(..., description="File where the side effect occurs") method: str = Field(..., description="Method where the side effect occurs") line: int = Field(..., description="Line number of the side effect") class UnresolvedDependency(BaseModel): """An unresolved dependency in the trace""" reason: str = Field(..., description="Reason why the dependency is unresolved") affected_file: str = Field(..., description="File affected by the unresolved dependency") line: int = Field(..., description="Line number of the unresolved dependency") class IncomingDependency(BaseModel): """An incoming dependency (what calls this target)""" from_file: str = Field(..., description="Source file of the dependency") from_class: str = Field(..., description="Source class of the dependency") from_method: str = Field(..., description="Source method of the dependency") line: int = Field(..., description="Line number of the dependency") type: str = Field(..., description="Type of dependency") class OutgoingDependency(BaseModel): """An outgoing dependency (what this target calls)""" to_file: str = Field(..., description="Target file of the dependency") to_class: str = Field(..., description="Target class of the dependency") to_method: str = Field(..., description="Target method of the dependency") line: int = Field(..., description="Line number of the dependency") type: str = Field(..., description="Type of dependency") class TypeDependency(BaseModel): """A type-level dependency (inheritance, imports, etc.)""" dependency_type: str = Field(..., description="Type of dependency") source_file: str = Field(..., description="Source file of the dependency") source_entity: str = Field(..., description="Source entity (class, module)") target: str = Field(..., description="Target entity") class StateAccess(BaseModel): """State access information""" file: str = Field(..., description="File where state is accessed") method: str = Field(..., description="Method accessing the state") access_type: str = Field(..., description="Type of access (reads, writes, etc.)") state_entity: str = Field(..., description="State entity being accessed") class TraceComplete(BaseModel): """Complete trace analysis response""" status: Literal["trace_complete"] = "trace_complete" trace_type: Literal["precision", "dependencies"] = Field(..., description="Type of trace performed") # Precision mode fields entry_point: Optional[TraceEntryPoint] = Field(None, description="Entry point for precision trace") call_path: Optional[list[CallPathStep]] = Field(default_factory=list, description="Call path for precision trace") branching_points: Optional[list[BranchingPoint]] = Field(default_factory=list, description="Branching points") side_effects: Optional[list[SideEffect]] = Field(default_factory=list, description="Side effects detected") unresolved: Optional[list[UnresolvedDependency]] = Field( default_factory=list, description="Unresolved dependencies" ) # Dependencies mode fields target: Optional[TraceTarget] = Field(None, description="Target for dependency analysis") incoming_dependencies: Optional[list[IncomingDependency]] = Field( default_factory=list, description="Incoming dependencies" ) outgoing_dependencies: Optional[list[OutgoingDependency]] = Field( default_factory=list, description="Outgoing dependencies" ) type_dependencies: Optional[list[TypeDependency]] = Field(default_factory=list, description="Type dependencies") state_access: Optional[list[StateAccess]] = Field(default_factory=list, description="State access information") class DiagnosticHypothesis(BaseModel): """A debugging hypothesis with context and next steps""" rank: int = Field(..., description="Ranking of this hypothesis (1 = most likely)") confidence: Literal["high", "medium", "low"] = Field(..., description="Confidence level") hypothesis: str = Field(..., description="Description of the potential root cause") reasoning: str = Field(..., description="Why this hypothesis is plausible") next_step: str = Field(..., description="Suggested action to test/validate this hypothesis") class StructuredDebugResponse(BaseModel): """Enhanced debug response with multiple hypotheses""" summary: str = Field(..., description="Brief summary of the issue") hypotheses: list[DiagnosticHypothesis] = Field(..., description="Ranked list of potential causes") immediate_actions: list[str] = Field( default_factory=list, description="Immediate steps to take regardless of root cause", ) additional_context_needed: Optional[list[str]] = Field( default_factory=list, description="Additional files or information that would help with analysis", ) class DebugHypothesis(BaseModel): """A debugging hypothesis with detailed analysis""" name: str = Field(..., description="Name/title of the hypothesis") confidence: Literal["High", "Medium", "Low"] = Field(..., description="Confidence level") root_cause: str = Field(..., description="Technical explanation of the root cause") evidence: str = Field(..., description="Logs or code clues supporting this hypothesis") correlation: str = Field(..., description="How symptoms map to the cause") validation: str = Field(..., description="Quick test to confirm the hypothesis") minimal_fix: str = Field(..., description="Smallest change to resolve the issue") regression_check: str = Field(..., description="Why this fix is safe") file_references: list[str] = Field(default_factory=list, description="File:line format for exact locations") class DebugAnalysisComplete(BaseModel): """Complete debugging analysis with systematic investigation tracking""" status: Literal["analysis_complete"] = "analysis_complete" investigation_id: str = Field(..., description="Auto-generated unique ID for this investigation") summary: str = Field(..., description="Brief description of the problem and its impact") investigation_steps: list[str] = Field(..., description="Steps taken during the investigation") hypotheses: list[DebugHypothesis] = Field(..., description="Ranked hypotheses with detailed analysis") key_findings: list[str] = Field(..., description="Important discoveries made during analysis") immediate_actions: list[str] = Field(..., description="Steps to take regardless of which hypothesis is correct") recommended_tools: list[str] = Field(default_factory=list, description="Additional tools recommended for analysis") prevention_strategy: Optional[str] = Field( None, description="Targeted measures to prevent this exact issue from recurring" ) investigation_summary: str = Field( ..., description="Comprehensive summary of the complete investigation process and conclusions" ) class NoBugFound(BaseModel): """Response when thorough investigation finds no concrete evidence of a bug""" status: Literal["no_bug_found"] = "no_bug_found" summary: str = Field(..., description="Summary of what was thoroughly investigated") investigation_steps: list[str] = Field(..., description="Steps taken during the investigation") areas_examined: list[str] = Field(..., description="Code areas and potential failure points examined") confidence_level: Literal["High", "Medium", "Low"] = Field( ..., description="Confidence level in the no-bug finding" ) alternative_explanations: list[str] = Field( ..., description="Possible alternative explanations for reported symptoms" ) recommended_questions: list[str] = Field(..., description="Questions to clarify the issue with the user") next_steps: list[str] = Field(..., description="Suggested actions to better understand the reported issue") # Registry mapping status strings to their corresponding Pydantic models SPECIAL_STATUS_MODELS = { "files_required_to_continue": FilesNeededRequest, "full_codereview_required": FullCodereviewRequired, "focused_review_required": FocusedReviewRequired, "test_sample_needed": TestSampleNeeded, "more_tests_required": MoreTestsRequired, "refactor_analysis_complete": RefactorAnalysisComplete, "trace_complete": TraceComplete, "resend_prompt": ResendPromptRequest, "code_too_large": CodeTooLargeRequest, "analysis_complete": DebugAnalysisComplete, "no_bug_found": NoBugFound, }