563 lines
23 KiB
Python
563 lines
23 KiB
Python
"""
|
|
Pydantic AI powered coding runtime with Claude Code style tools.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import shlex
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
from pydantic_ai import Agent, RunContext
|
|
|
|
from .config import (
|
|
AgentMetadata,
|
|
AzureBlobResourceConfig,
|
|
CodingRequestConfig,
|
|
DatabaseEngine,
|
|
DatabaseResourceConfig,
|
|
GitResourceConfig,
|
|
LiteLLMConfig,
|
|
ResourceConfig,
|
|
WorkspaceConfig,
|
|
)
|
|
from .resources import build_authenticated_repo_url, safe_workspace_path, summarize_resources
|
|
|
|
|
|
class ToolEvent(BaseModel):
|
|
tool: str
|
|
payload: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class CodingRunContext(BaseModel):
|
|
workspace: WorkspaceConfig
|
|
resources: ResourceConfig
|
|
changed_files: list[str] = Field(default_factory=list)
|
|
tool_log: list[ToolEvent] = Field(default_factory=list)
|
|
finish_summary: Optional[str] = None
|
|
|
|
model_config = {"arbitrary_types_allowed": True}
|
|
|
|
|
|
class CodingRunResult(BaseModel):
|
|
response_text: str
|
|
summary: str
|
|
workspace_root: str
|
|
files_changed: list[str] = Field(default_factory=list)
|
|
tool_log: list[ToolEvent] = Field(default_factory=list)
|
|
resources_used: dict[str, Any] = Field(default_factory=dict)
|
|
|
|
|
|
class CodingRuntimeError(Exception):
|
|
"""Structured runtime error that can be returned through JSON-RPC."""
|
|
|
|
def __init__(
|
|
self,
|
|
message: str,
|
|
*,
|
|
code: str = "runtime_error",
|
|
stage: str = "runtime",
|
|
data: Optional[dict[str, Any]] = None,
|
|
):
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.stage = stage
|
|
self.data = data or {}
|
|
|
|
def to_payload(self) -> dict[str, Any]:
|
|
payload = {
|
|
"code": self.code,
|
|
"stage": self.stage,
|
|
}
|
|
payload.update(self.data)
|
|
return payload
|
|
|
|
|
|
class WorkspacePreparationError(CodingRuntimeError):
|
|
"""Raised when the remote workspace cannot be prepared safely."""
|
|
|
|
|
|
class CodingA2ARuntime:
|
|
def __init__(
|
|
self,
|
|
llm_config: LiteLLMConfig,
|
|
metadata: AgentMetadata,
|
|
):
|
|
self.llm_config = llm_config
|
|
self.metadata = metadata
|
|
self._apply_llm_env(self.llm_config)
|
|
self._agent = self._build_agent(self.llm_config)
|
|
|
|
def _build_agent(self, llm_config: LiteLLMConfig) -> Agent:
|
|
agent: Agent[CodingRunContext] = Agent(
|
|
llm_config.normalized_model,
|
|
system_prompt=self.metadata.effective_system_prompt,
|
|
deps_type=CodingRunContext,
|
|
)
|
|
|
|
@agent.tool
|
|
async def read_file(ctx: RunContext[CodingRunContext], path: str) -> str:
|
|
target = safe_workspace_path(
|
|
ctx.deps.workspace.root_dir,
|
|
path,
|
|
ctx.deps.workspace.allowed_paths or ctx.deps.resources.git.allowed_paths if ctx.deps.resources.git else ctx.deps.workspace.allowed_paths,
|
|
)
|
|
content = target.read_text(encoding="utf-8", errors="replace")
|
|
ctx.deps.tool_log.append(ToolEvent(tool="read_file", payload={"path": path, "bytes": len(content)}))
|
|
return content
|
|
|
|
@agent.tool
|
|
async def list_files(ctx: RunContext[CodingRunContext], glob_pattern: str = "**/*") -> str:
|
|
root = Path(ctx.deps.workspace.root_dir).resolve()
|
|
matched = [
|
|
path.relative_to(root).as_posix()
|
|
for path in sorted(root.glob(glob_pattern))
|
|
if path.is_file()
|
|
]
|
|
ctx.deps.tool_log.append(ToolEvent(tool="list_files", payload={"pattern": glob_pattern, "count": len(matched)}))
|
|
return "\n".join(matched[:500]) if matched else "(no files matched)"
|
|
|
|
@agent.tool
|
|
async def write_file(ctx: RunContext[CodingRunContext], path: str, content: str) -> str:
|
|
target = safe_workspace_path(
|
|
ctx.deps.workspace.root_dir,
|
|
path,
|
|
ctx.deps.workspace.allowed_paths or ctx.deps.resources.git.allowed_paths if ctx.deps.resources.git else ctx.deps.workspace.allowed_paths,
|
|
)
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
target.write_text(content, encoding="utf-8")
|
|
rel = target.relative_to(Path(ctx.deps.workspace.root_dir).resolve()).as_posix()
|
|
if rel not in ctx.deps.changed_files:
|
|
ctx.deps.changed_files.append(rel)
|
|
ctx.deps.tool_log.append(ToolEvent(tool="write_file", payload={"path": rel, "bytes": len(content.encode())}))
|
|
return f"written {rel}"
|
|
|
|
@agent.tool
|
|
async def edit_file(ctx: RunContext[CodingRunContext], path: str, old_text: str, new_text: str) -> str:
|
|
target = safe_workspace_path(
|
|
ctx.deps.workspace.root_dir,
|
|
path,
|
|
ctx.deps.workspace.allowed_paths or ctx.deps.resources.git.allowed_paths if ctx.deps.resources.git else ctx.deps.workspace.allowed_paths,
|
|
)
|
|
content = target.read_text(encoding="utf-8", errors="replace")
|
|
if old_text not in content:
|
|
return "old_text not found"
|
|
updated = content.replace(old_text, new_text, 1)
|
|
target.write_text(updated, encoding="utf-8")
|
|
rel = target.relative_to(Path(ctx.deps.workspace.root_dir).resolve()).as_posix()
|
|
if rel not in ctx.deps.changed_files:
|
|
ctx.deps.changed_files.append(rel)
|
|
ctx.deps.tool_log.append(ToolEvent(tool="edit_file", payload={"path": rel}))
|
|
return f"edited {rel}"
|
|
|
|
@agent.tool
|
|
async def run_command(ctx: RunContext[CodingRunContext], command: str) -> str:
|
|
completed = subprocess.run(
|
|
["bash", "-lc", command],
|
|
cwd=ctx.deps.workspace.root_dir,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300,
|
|
)
|
|
ctx.deps.tool_log.append(
|
|
ToolEvent(
|
|
tool="run_command",
|
|
payload={"command": command, "returncode": completed.returncode},
|
|
)
|
|
)
|
|
output = completed.stdout.strip()
|
|
if completed.stderr.strip():
|
|
output = f"{output}\n{completed.stderr.strip()}".strip()
|
|
return output or f"(exit code {completed.returncode})"
|
|
|
|
@agent.tool
|
|
async def git_prepare_workspace(ctx: RunContext[CodingRunContext]) -> str:
|
|
git = ctx.deps.resources.git
|
|
if not git or not git.repo_url:
|
|
return "git resource not configured"
|
|
workspace = Path(ctx.deps.workspace.root_dir)
|
|
if (workspace / ".git").exists():
|
|
return "workspace already contains a git repository"
|
|
workspace.parent.mkdir(parents=True, exist_ok=True)
|
|
auth_url = build_authenticated_repo_url(
|
|
git.repo_url,
|
|
username=git.username,
|
|
password=git.password,
|
|
token=git.token,
|
|
)
|
|
branch = git.default_branch or "main"
|
|
completed = subprocess.run(
|
|
["git", "clone", "--branch", branch, auth_url, workspace.name],
|
|
cwd=str(workspace.parent),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300,
|
|
)
|
|
ctx.deps.tool_log.append(
|
|
ToolEvent(
|
|
tool="git_prepare_workspace",
|
|
payload={"repo_url": git.repo_url, "returncode": completed.returncode},
|
|
)
|
|
)
|
|
if completed.returncode != 0:
|
|
return (completed.stdout + "\n" + completed.stderr).strip()
|
|
return f"cloned {git.repo_url} into {workspace}"
|
|
|
|
@agent.tool
|
|
async def git_status(ctx: RunContext[CodingRunContext]) -> str:
|
|
return await run_command(ctx, "git status --short")
|
|
|
|
@agent.tool
|
|
async def git_diff(ctx: RunContext[CodingRunContext], ref: str = "HEAD") -> str:
|
|
return await run_command(ctx, f"git diff {ref}")
|
|
|
|
@agent.tool
|
|
async def git_create_branch(ctx: RunContext[CodingRunContext], branch_name: str) -> str:
|
|
return await run_command(ctx, f"git checkout -B {branch_name}")
|
|
|
|
@agent.tool
|
|
async def git_commit(ctx: RunContext[CodingRunContext], message: str) -> str:
|
|
await run_command(ctx, "git add -A")
|
|
return await run_command(ctx, f"git commit -m {shlex.quote(message)}")
|
|
|
|
@agent.tool
|
|
async def git_push(ctx: RunContext[CodingRunContext], remote: str = "origin", branch_name: Optional[str] = None) -> str:
|
|
git = ctx.deps.resources.git
|
|
if git and git.repo_url:
|
|
auth_url = build_authenticated_repo_url(
|
|
git.repo_url,
|
|
username=git.username,
|
|
password=git.password,
|
|
token=git.token,
|
|
)
|
|
subprocess.run(
|
|
["git", "remote", "set-url", remote, auth_url],
|
|
cwd=ctx.deps.workspace.root_dir,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
)
|
|
target = branch_name or "HEAD"
|
|
return await run_command(ctx, f"git push {remote} {target}")
|
|
|
|
@agent.tool
|
|
async def list_database_tables(ctx: RunContext[CodingRunContext], engine: str = "postgresql") -> str:
|
|
config = self._select_database_config(ctx.deps.resources, engine)
|
|
if not config:
|
|
return f"{engine} resource not configured"
|
|
result = self._run_database_query(config, self._default_table_query(config.engine))
|
|
ctx.deps.tool_log.append(ToolEvent(tool="list_database_tables", payload={"engine": config.engine.value}))
|
|
return result
|
|
|
|
@agent.tool
|
|
async def run_database_query(ctx: RunContext[CodingRunContext], engine: str, query: str) -> str:
|
|
config = self._select_database_config(ctx.deps.resources, engine)
|
|
if not config:
|
|
return f"{engine} resource not configured"
|
|
result = self._run_database_query(config, query)
|
|
ctx.deps.tool_log.append(ToolEvent(tool="run_database_query", payload={"engine": config.engine.value}))
|
|
return result
|
|
|
|
@agent.tool
|
|
async def list_blob_objects(ctx: RunContext[CodingRunContext], limit: int = 50) -> str:
|
|
config = ctx.deps.resources.azure_blob
|
|
if not config:
|
|
return "azure_blob resource not configured"
|
|
result = self._list_blob_objects(config, limit=limit)
|
|
ctx.deps.tool_log.append(ToolEvent(tool="list_blob_objects", payload={"count": len(result)}))
|
|
return json.dumps(result, ensure_ascii=False, indent=2)
|
|
|
|
@agent.tool
|
|
async def read_blob_text(ctx: RunContext[CodingRunContext], blob_name: str, encoding: str = "utf-8") -> str:
|
|
config = ctx.deps.resources.azure_blob
|
|
if not config:
|
|
return "azure_blob resource not configured"
|
|
text = self._read_blob_text(config, blob_name=blob_name, encoding=encoding)
|
|
ctx.deps.tool_log.append(ToolEvent(tool="read_blob_text", payload={"blob_name": blob_name, "bytes": len(text.encode())}))
|
|
return text
|
|
|
|
@agent.tool
|
|
async def finish(ctx: RunContext[CodingRunContext], summary: str) -> str:
|
|
ctx.deps.finish_summary = summary
|
|
ctx.deps.tool_log.append(ToolEvent(tool="finish", payload={"summary": summary}))
|
|
return f"done: {summary}"
|
|
|
|
return agent
|
|
|
|
async def run_task(
|
|
self,
|
|
prompt: str,
|
|
request_config: CodingRequestConfig,
|
|
*,
|
|
api_key: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
) -> CodingRunResult:
|
|
llm_config = LiteLLMConfig(
|
|
base_url=self.llm_config.base_url,
|
|
api_key=api_key or request_config.api_key or self.llm_config.api_key,
|
|
model=model or request_config.model or self.llm_config.model,
|
|
timeout=self.llm_config.timeout,
|
|
max_tokens=self.llm_config.max_tokens,
|
|
)
|
|
previous_api_key = os.environ.get("OPENAI_API_KEY")
|
|
previous_base_url = os.environ.get("OPENAI_BASE_URL")
|
|
self._apply_llm_env(llm_config)
|
|
|
|
agent = self._build_agent(llm_config)
|
|
deps = CodingRunContext(workspace=request_config.workspace, resources=request_config.resources)
|
|
self._prepare_workspace(request_config, deps)
|
|
|
|
initial_prompt = self._build_initial_prompt(prompt, request_config)
|
|
|
|
try:
|
|
result = await agent.run(initial_prompt, deps=deps)
|
|
response_text = getattr(result, "output", None) or getattr(result, "data", None) or str(result)
|
|
summary = deps.finish_summary or response_text
|
|
return CodingRunResult(
|
|
response_text=str(response_text),
|
|
summary=summary,
|
|
workspace_root=request_config.workspace.root_dir,
|
|
files_changed=deps.changed_files,
|
|
tool_log=deps.tool_log,
|
|
resources_used=summarize_resources(request_config.resources.model_dump(exclude_none=True)),
|
|
)
|
|
finally:
|
|
if previous_api_key is None:
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
else:
|
|
os.environ["OPENAI_API_KEY"] = previous_api_key
|
|
if previous_base_url is None:
|
|
os.environ.pop("OPENAI_BASE_URL", None)
|
|
else:
|
|
os.environ["OPENAI_BASE_URL"] = previous_base_url
|
|
|
|
def _apply_llm_env(self, llm_config: LiteLLMConfig) -> None:
|
|
if llm_config.api_key:
|
|
os.environ["OPENAI_API_KEY"] = llm_config.api_key
|
|
os.environ["OPENAI_BASE_URL"] = llm_config.base_url
|
|
|
|
def _prepare_workspace(
|
|
self,
|
|
request_config: CodingRequestConfig,
|
|
deps: CodingRunContext,
|
|
) -> None:
|
|
workspace_root = Path(request_config.workspace.root_dir).expanduser().resolve()
|
|
request_config.workspace.root_dir = str(workspace_root)
|
|
deps.workspace.root_dir = str(workspace_root)
|
|
|
|
if workspace_root.exists() and not workspace_root.is_dir():
|
|
raise WorkspacePreparationError(
|
|
"workspace root is not a directory",
|
|
code="workspace_not_directory",
|
|
stage="workspace_prepare",
|
|
data={"workspace_root": str(workspace_root)},
|
|
)
|
|
|
|
git = request_config.resources.git
|
|
if git and git.repo_url:
|
|
self._prepare_git_workspace(workspace_root, git, deps)
|
|
return
|
|
|
|
workspace_root.mkdir(parents=True, exist_ok=True)
|
|
deps.tool_log.append(
|
|
ToolEvent(
|
|
tool="workspace_prepare",
|
|
payload={"workspace_root": str(workspace_root), "mode": "empty_workspace"},
|
|
)
|
|
)
|
|
|
|
def _prepare_git_workspace(
|
|
self,
|
|
workspace_root: Path,
|
|
git: GitResourceConfig,
|
|
deps: CodingRunContext,
|
|
) -> None:
|
|
if (workspace_root / ".git").exists():
|
|
deps.tool_log.append(
|
|
ToolEvent(
|
|
tool="git_prepare_workspace",
|
|
payload={
|
|
"workspace_root": str(workspace_root),
|
|
"repo_url": git.repo_url,
|
|
"mode": "existing_repository",
|
|
},
|
|
)
|
|
)
|
|
return
|
|
|
|
workspace_root.parent.mkdir(parents=True, exist_ok=True)
|
|
if workspace_root.exists():
|
|
if not workspace_root.is_dir():
|
|
raise WorkspacePreparationError(
|
|
"workspace root is not a directory",
|
|
code="workspace_not_directory",
|
|
stage="git_prepare_workspace",
|
|
data={"workspace_root": str(workspace_root), "repo_url": git.repo_url},
|
|
)
|
|
if any(workspace_root.iterdir()):
|
|
raise WorkspacePreparationError(
|
|
"workspace already exists but is not a git repository",
|
|
code="workspace_not_git_repository",
|
|
stage="git_prepare_workspace",
|
|
data={"workspace_root": str(workspace_root), "repo_url": git.repo_url},
|
|
)
|
|
workspace_root.rmdir()
|
|
|
|
auth_url = build_authenticated_repo_url(
|
|
git.repo_url,
|
|
username=git.username,
|
|
password=git.password,
|
|
token=git.token,
|
|
)
|
|
branch = git.default_branch or "main"
|
|
try:
|
|
completed = subprocess.run(
|
|
["git", "clone", "--branch", branch, auth_url, workspace_root.name],
|
|
cwd=str(workspace_root.parent),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=180,
|
|
)
|
|
except subprocess.TimeoutExpired as exc:
|
|
raise WorkspacePreparationError(
|
|
"git workspace preparation timed out",
|
|
code="git_prepare_timeout",
|
|
stage="git_prepare_workspace",
|
|
data={
|
|
"workspace_root": str(workspace_root),
|
|
"repo_url": git.repo_url,
|
|
"branch": branch,
|
|
"timeout_seconds": exc.timeout,
|
|
},
|
|
) from exc
|
|
|
|
payload = {
|
|
"workspace_root": str(workspace_root),
|
|
"repo_url": git.repo_url,
|
|
"branch": branch,
|
|
"returncode": completed.returncode,
|
|
}
|
|
deps.tool_log.append(ToolEvent(tool="git_prepare_workspace", payload=payload))
|
|
|
|
if completed.returncode != 0:
|
|
raise WorkspacePreparationError(
|
|
"git workspace preparation failed",
|
|
code="git_prepare_failed",
|
|
stage="git_prepare_workspace",
|
|
data={
|
|
**payload,
|
|
"stdout": completed.stdout.strip(),
|
|
"stderr": completed.stderr.strip(),
|
|
},
|
|
)
|
|
|
|
def _build_initial_prompt(self, prompt: str, request_config: CodingRequestConfig) -> str:
|
|
workspace = request_config.workspace
|
|
parts = [
|
|
f"USER TASK:\n{prompt}",
|
|
f"WORKSPACE ROOT: {workspace.root_dir}",
|
|
f"TASK MODE: {request_config.task_mode}",
|
|
]
|
|
if workspace.entry_file:
|
|
parts.append(f"START BY READING: {workspace.entry_file}")
|
|
if workspace.context_files:
|
|
parts.append("ALSO CONSIDER: " + ", ".join(workspace.context_files))
|
|
if workspace.allowed_paths:
|
|
parts.append("YOU MAY ONLY MODIFY: " + ", ".join(workspace.allowed_paths))
|
|
if request_config.branch_name:
|
|
parts.append(f"PREFERRED BRANCH: {request_config.branch_name}")
|
|
if request_config.commit_message:
|
|
parts.append(f"SUGGESTED COMMIT MESSAGE: {request_config.commit_message}")
|
|
parts.append(
|
|
"Use tools to inspect before editing. Prefer minimal precise changes. "
|
|
"When the work is complete, call finish(summary)."
|
|
)
|
|
return "\n\n".join(parts)
|
|
|
|
def _select_database_config(self, resources: ResourceConfig, engine: str) -> Optional[DatabaseResourceConfig]:
|
|
requested = engine.lower()
|
|
if requested in {"mysql", "mariadb"}:
|
|
return resources.mysql
|
|
return resources.postgresql
|
|
|
|
def _default_table_query(self, engine: DatabaseEngine) -> str:
|
|
if engine == DatabaseEngine.mysql:
|
|
return "SHOW TABLES"
|
|
return (
|
|
"SELECT table_schema, table_name FROM information_schema.tables "
|
|
"WHERE table_schema NOT IN ('pg_catalog', 'information_schema') "
|
|
"ORDER BY table_schema, table_name LIMIT 200"
|
|
)
|
|
|
|
def _run_database_query(self, config: DatabaseResourceConfig, query: str) -> str:
|
|
if config.engine == DatabaseEngine.mysql:
|
|
import pymysql
|
|
|
|
connection = pymysql.connect(
|
|
host=config.host,
|
|
port=config.port or 3306,
|
|
user=config.username,
|
|
password=config.password,
|
|
database=config.database,
|
|
cursorclass=pymysql.cursors.DictCursor,
|
|
connect_timeout=10,
|
|
)
|
|
else:
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
|
|
connection = psycopg2.connect(
|
|
host=config.host,
|
|
port=config.port or 5432,
|
|
user=config.username,
|
|
password=config.password,
|
|
dbname=config.database,
|
|
connect_timeout=10,
|
|
sslmode=config.ssl_mode or "prefer",
|
|
)
|
|
try:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute(query)
|
|
rows = cursor.fetchall()
|
|
return json.dumps(rows, ensure_ascii=False, default=str, indent=2)
|
|
finally:
|
|
connection.close()
|
|
|
|
def _get_blob_client(self, config: AzureBlobResourceConfig):
|
|
from azure.storage.blob import BlobServiceClient
|
|
|
|
if config.connection_string:
|
|
return BlobServiceClient.from_connection_string(config.connection_string)
|
|
if config.account_url and config.sas_token:
|
|
return BlobServiceClient(account_url=config.account_url, credential=config.sas_token)
|
|
if config.account_url and config.account_key:
|
|
return BlobServiceClient(account_url=config.account_url, credential=config.account_key)
|
|
if config.account_name and config.account_key:
|
|
account_url = f"https://{config.account_name}.blob.core.windows.net"
|
|
return BlobServiceClient(account_url=account_url, credential=config.account_key)
|
|
raise ValueError("azure blob credentials are not configured")
|
|
|
|
def _list_blob_objects(self, config: AzureBlobResourceConfig, limit: int = 50) -> list[dict[str, Any]]:
|
|
service = self._get_blob_client(config)
|
|
container = service.get_container_client(config.container_name)
|
|
items = []
|
|
for index, blob in enumerate(container.list_blobs(name_starts_with=config.prefix or None)):
|
|
if index >= limit:
|
|
break
|
|
items.append(
|
|
{
|
|
"name": blob.name,
|
|
"size": blob.size,
|
|
"content_type": getattr(blob.content_settings, "content_type", None),
|
|
}
|
|
)
|
|
return items
|
|
|
|
def _read_blob_text(self, config: AzureBlobResourceConfig, blob_name: str, encoding: str = "utf-8") -> str:
|
|
service = self._get_blob_client(config)
|
|
blob_client = service.get_blob_client(container=config.container_name, blob=blob_name)
|
|
return blob_client.download_blob().readall().decode(encoding, errors="replace")
|