chore(agent): pydantic models

Author: bgagent

agent/src/config.py

@@ -4,7 +4,7 @@
 import sys
 import uuid
 
-from models import TaskType
+from models import TaskConfig, TaskType
 
 AGENT_WORKSPACE = os.environ.get("AGENT_WORKSPACE", "/workspace")
 
@@ -52,54 +52,60 @@ def build_config(
     task_type: str = "new_task",
     branch_name: str = "",
     pr_number: str = "",
-) -> dict:
+) -> TaskConfig:
     """Build and validate configuration from explicit parameters.
 
     Parameters fall back to environment variables if empty.
     """
-    config = {
-        "repo_url": repo_url or os.environ.get("REPO_URL", ""),
-        "issue_number": issue_number or os.environ.get("ISSUE_NUMBER", ""),
-        "task_description": task_description or os.environ.get("TASK_DESCRIPTION", ""),
-        "github_token": github_token or resolve_github_token(),
-        "aws_region": aws_region or os.environ.get("AWS_REGION", ""),
-        "anthropic_model": anthropic_model
-        or os.environ.get("ANTHROPIC_MODEL", "us.anthropic.claude-sonnet-4-6"),
-        "dry_run": dry_run,
-        "max_turns": max_turns,
-        "max_budget_usd": max_budget_usd,
-        "system_prompt_overrides": system_prompt_overrides,
-        "task_type": task_type,
-        "branch_name": branch_name,
-        "pr_number": pr_number,
-    }
+    resolved_repo_url = repo_url or os.environ.get("REPO_URL", "")
+    resolved_issue_number = issue_number or os.environ.get("ISSUE_NUMBER", "")
+    resolved_task_description = task_description or os.environ.get("TASK_DESCRIPTION", "")
+    resolved_github_token = github_token or resolve_github_token()
+    resolved_aws_region = aws_region or os.environ.get("AWS_REGION", "")
+    resolved_anthropic_model = anthropic_model or os.environ.get(
+        "ANTHROPIC_MODEL", "us.anthropic.claude-sonnet-4-6"
+    )
 
     errors = []
-    if not config["repo_url"]:
+    if not resolved_repo_url:
         errors.append("repo_url is required (e.g., 'owner/repo')")
-    if not config["github_token"]:
+    if not resolved_github_token:
         errors.append("github_token is required")
-    if not config["aws_region"]:
+    if not resolved_aws_region:
         errors.append("aws_region is required for Bedrock")
     try:
-        task = TaskType(config["task_type"])
+        task = TaskType(task_type)
     except ValueError:
-        errors.append(f"Invalid task_type: '{config['task_type']}'")
+        errors.append(f"Invalid task_type: '{task_type}'")
         task = None
     if task and task.is_pr_task:
-        if not config["pr_number"]:
+        if not pr_number:
             errors.append("pr_number is required for pr_iteration/pr_review task type")
-    elif task and not config["issue_number"] and not config["task_description"]:
+    elif task and not resolved_issue_number and not resolved_task_description:
         errors.append("Either issue_number or task_description is required")
 
     if errors:
         raise ValueError("; ".join(errors))
 
-    config["task_id"] = task_id or uuid.uuid4().hex[:12]
-    return config
-
-
-def get_config() -> dict:
+    return TaskConfig(
+        repo_url=resolved_repo_url,
+        issue_number=resolved_issue_number,
+        task_description=resolved_task_description,
+        github_token=resolved_github_token,
+        aws_region=resolved_aws_region,
+        anthropic_model=resolved_anthropic_model,
+        dry_run=dry_run,
+        max_turns=max_turns,
+        max_budget_usd=max_budget_usd,
+        system_prompt_overrides=system_prompt_overrides,
+        task_type=task_type,
+        branch_name=branch_name,
+        pr_number=pr_number,
+        task_id=task_id or uuid.uuid4().hex[:12],
+    )
+
+
+def get_config() -> TaskConfig:
     """Parse configuration from environment variables (local batch mode)."""
     try:
         return build_config(

agent/src/context.py

@@ -2,8 +2,10 @@
 
 import requests
 
+from models import GitHubIssue, IssueComment, TaskConfig
 
-def fetch_github_issue(repo_url: str, issue_number: str, token: str) -> dict:
+
+def fetch_github_issue(repo_url: str, issue_number: str, token: str) -> GitHubIssue:
     """Fetch a GitHub issue's title, body, and comments."""
     headers = {
         "Authorization": f"token {token}",
@@ -20,51 +22,54 @@ def fetch_github_issue(repo_url: str, issue_number: str, token: str) -> dict:
     issue = issue_resp.json()
 
     # Fetch comments
-    comments = []
+    comments: list[IssueComment] = []
     if issue.get("comments", 0) > 0:
         comments_resp = requests.get(
             f"https://api.github.com/repos/{repo_url}/issues/{issue_number}/comments",
             headers=headers,
             timeout=30,
         )
         comments_resp.raise_for_status()
-        comments = [{"author": c["user"]["login"], "body": c["body"]} for c in comments_resp.json()]
+        comments = [
+            IssueComment(author=c["user"]["login"], body=c["body"]) for c in comments_resp.json()
+        ]
 
-    return {
-        "title": issue["title"],
-        "body": issue.get("body", ""),
-        "number": issue["number"],
-        "comments": comments,
-    }
+    return GitHubIssue(
+        title=issue["title"],
+        body=issue.get("body", "") or "",
+        number=issue["number"],
+        comments=comments,
+    )
 
 
-def assemble_prompt(config: dict) -> str:
+def assemble_prompt(config: TaskConfig) -> str:
     """Assemble the user prompt from issue context and task description.
 
     .. deprecated::
         In production (AgentCore server mode), the orchestrator's
         ``assembleUserPrompt()`` in ``context-hydration.ts`` is the sole prompt
-        assembler. The hydrated prompt arrives via ``hydrated_context["user_prompt"]``.
+        assembler. The hydrated prompt arrives via
+        ``HydratedContext.user_prompt`` (validated from the incoming JSON).
         This Python implementation is retained only for **local batch mode**
         (``python src/entrypoint.py``) and **dry-run mode** (``DRY_RUN=1``).
     """
     parts = []
 
-    parts.append(f"Task ID: {config['task_id']}")
-    parts.append(f"Repository: {config['repo_url']}")
+    parts.append(f"Task ID: {config.task_id}")
+    parts.append(f"Repository: {config.repo_url}")
 
-    if config.get("issue"):
-        issue = config["issue"]
-        parts.append(f"\n## GitHub Issue #{issue['number']}: {issue['title']}\n")
-        parts.append(issue["body"] or "(no description)")
-        if issue["comments"]:
+    if config.issue:
+        issue = config.issue
+        parts.append(f"\n## GitHub Issue #{issue.number}: {issue.title}\n")
+        parts.append(issue.body or "(no description)")
+        if issue.comments:
             parts.append("\n### Comments\n")
-            for c in issue["comments"]:
-                parts.append(f"**@{c['author']}**: {c['body']}\n")
+            for c in issue.comments:
+                parts.append(f"**@{c.author}**: {c.body}\n")
 
-    if config["task_description"]:
-        parts.append(f"\n## Task\n\n{config['task_description']}")
-    elif config.get("issue"):
+    if config.task_description:
+        parts.append(f"\n## Task\n\n{config.task_description}")
+    elif config.issue:
         parts.append(
             "\n## Task\n\nResolve the GitHub issue described above. "
             "Follow the workflow in your system instructions."

agent/src/entrypoint.py

@@ -20,6 +20,18 @@
     resolve_github_token,
 )
 from context import assemble_prompt, fetch_github_issue  # noqa: F401
+from models import (  # noqa: F401
+    AgentResult,
+    GitHubIssue,
+    HydratedContext,
+    IssueComment,
+    MemoryContext,
+    RepoSetup,
+    TaskConfig,
+    TaskResult,
+    TaskType,
+    TokenUsage,
+)
 from pipeline import main, run_task  # noqa: F401
 from post_hooks import (  # noqa: F401
     ensure_committed,

agent/src/models.py

@@ -1,7 +1,11 @@
 """Data models and enumerations for the agent pipeline."""
 
+from __future__ import annotations
+
 from enum import StrEnum
 
+from pydantic import BaseModel, ConfigDict
+
 
 class TaskType(StrEnum):
     """Supported task types."""
@@ -17,3 +21,114 @@ def is_pr_task(self) -> bool:
     @property
     def is_read_only(self) -> bool:
         return self == TaskType.pr_review
+
+
+class IssueComment(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    author: str
+    body: str
+
+
+class GitHubIssue(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    title: str
+    body: str = ""
+    number: int
+    comments: list[IssueComment] = []
+
+
+class MemoryContext(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    repo_knowledge: list[str] = []
+    past_episodes: list[str] = []
+
+
+class HydratedContext(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    user_prompt: str
+    issue: GitHubIssue | None = None
+    resolved_base_branch: str | None = None
+    truncated: bool = False
+    memory_context: MemoryContext | None = None
+
+
+class TaskConfig(BaseModel):
+    model_config = ConfigDict(validate_assignment=True)
+
+    repo_url: str
+    issue_number: str = ""
+    task_description: str = ""
+    github_token: str
+    aws_region: str
+    anthropic_model: str = "us.anthropic.claude-sonnet-4-6"
+    dry_run: bool = False
+    max_turns: int = 10
+    max_budget_usd: float | None = None
+    system_prompt_overrides: str = ""
+    task_type: str = "new_task"
+    branch_name: str = ""
+    pr_number: str = ""
+    task_id: str = ""
+    # Enriched mid-flight by pipeline.py:
+    cedar_policies: list[str] = []
+    issue: GitHubIssue | None = None
+    base_branch: str | None = None
+
+
+class RepoSetup(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    repo_dir: str
+    branch: str
+    notes: list[str] = []
+    build_before: bool = True
+    lint_before: bool = True
+    default_branch: str = "main"
+
+
+class TokenUsage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read_input_tokens: int = 0
+    cache_creation_input_tokens: int = 0
+
+
+class AgentResult(BaseModel):
+    status: str = "unknown"
+    turns: int = 0
+    num_turns: int = 0
+    cost_usd: float | None = None
+    duration_ms: int = 0
+    duration_api_ms: int = 0
+    session_id: str = ""
+    error: str | None = None
+    usage: TokenUsage | None = None
+
+
+class TaskResult(BaseModel):
+    status: str
+    agent_status: str = "unknown"
+    pr_url: str | None = None
+    build_passed: bool = False
+    lint_passed: bool = False
+    cost_usd: float | None = None
+    turns: int | None = None
+    duration_s: float = 0.0
+    task_id: str = ""
+    disk_before: str = ""
+    disk_after: str = ""
+    disk_delta: str = ""
+    prompt_version: str | None = None
+    memory_written: bool = False
+    error: str | None = None
+    session_id: str | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cache_read_input_tokens: int | None = None
+    cache_creation_input_tokens: int | None = None