"""Job handler dispatch + default handlers (AI + logging). Handlers share the following signature:: async fn( envelope: JobRequestEnvelope, *, deps: WorkerDependencies, ) -> HandlerResult WorkerDependencies wraps Mongo collections, the AI Gateway, and logging. """ from __future__ import annotations from dataclasses import dataclass, field from typing import Any, Awaitable, Callable, Protocol from uuid import UUID import structlog from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase from asf_events.base import JobRequestEnvelope from asf_domain import JobStatus, JobUpdate, JobUsage from app.config import WorkerSettings # type: ignore[attr-defined] log = structlog.get_logger(__name__) JobHandlerFn = Callable[["JobRequestEnvelope", "WorkerDependencies"], Awaitable["HandlerResult"]] @dataclass class HandlerResult: status: JobStatus result: dict[str, Any] = field(default_factory=dict) error: str | None = None error_detail: dict[str, Any] | None = None usage: JobUsage | None = None @dataclass class WorkerDependencies: db: AsyncIOMotorDatabase jobs: AsyncIOMotorCollection ai_executions: AsyncIOMotorCollection artifacts: AsyncIOMotorCollection settings: WorkerSettings ai_gateway: Any | None = None # will be set if enable_ai_handlers=True worker_id: str = "asf-worker" # ========================================================================== # Default handlers # ========================================================================== async def _mark_job_done( deps: WorkerDependencies, job_id: UUID, result: HandlerResult, ) -> None: """Persist the final handler outcome back to MongoDB jobs collection. This is the ONLY place worker-state leaves process memory and hits the durable source of truth (MongoDB). """ updates: dict[str, Any] = { "status": result.status.value if isinstance(result.status, JobStatus) else str(result.status), "completed_at": __import__("datetime").datetime.now( __import__("datetime").timezone.utc ), "updated_at": __import__("datetime").datetime.now( __import__("datetime").timezone.utc ), "worker": deps.worker_id, } if result.result: updates["result"] = result.result if result.error: updates["error"] = result.error if result.error_detail: updates["error_detail"] = result.error_detail if result.usage is not None: updates["usage"] = ( result.usage.model_dump() if hasattr(result.usage, "model_dump") else result.usage ) await deps.jobs.update_one({"_id": job_id}, {"$set": updates}) async def noop_handler( env: JobRequestEnvelope, deps: WorkerDependencies ) -> HandlerResult: """Catch-all handler: acknowledges the job and leaves a note. Useful when a queue is not yet staffed with a real handler — ensures we never leave work stuck "running" with no outcome. """ log.warn( "no_handler_registered", job_id=str(env.job_id), queue=env.type, payload_keys=list(env.payload.keys()), ) return HandlerResult( status=JobStatus.FAILED, error=f"No worker handler registered for job_type={env.type!r}", error_detail={"hint": "register a handler in app.worker.HANDLERS_BY_TYPE"}, ) async def echo_handler( env: JobRequestEnvelope, deps: WorkerDependencies ) -> HandlerResult: """Simple echo / dev handler: stores the payload back as result.""" return HandlerResult( status=JobStatus.COMPLETED, result={"echo": env.payload, "handled_by": deps.worker_id}, usage=JobUsage(cpu_seconds=0.001), ) async def ai_handler( env: JobRequestEnvelope, deps: WorkerDependencies ) -> HandlerResult: """AI-interactive + AI-background handler. Dispatches to the ASF AI Gateway (MockProvider in bootstrap) for: - "ask_asf" kind (chat completion) - "analyze_code" - "generate_tests" - "decompose_requirement" In the bootstrap phase we safely fall back to echo if no gateway is attached (e.g. test runs with AI disabled). """ kind = env.payload.get("kind", "unknown") log.info( "ai_handler_start", job_id=str(env.job_id), kind=kind, provider=deps.settings.ai_gateway_provider, ) if deps.ai_gateway is None: return HandlerResult( status=JobStatus.COMPLETED, result={ "note": ( "ai_gateway not attached to this worker. " "Install services/asf-ai and wire it in worker." ), "payload": env.payload, }, ) from asf_ai import ( # type: ignore[import-not-found] AIGateway, ChatRequest, ChatMessage, TraceContext, CodeAnalysisRequest, TestGenerationRequest, RequirementDecomposeRequest, ) gw: AIGateway = deps.ai_gateway trace = TraceContext( tenant_id=env.tenant_id, project_id=env.project_id, job_id=env.job_id, correlation_id=env.correlation_id, ) try: if kind == "ask_asf": question = env.payload.get("question", "") req = ChatRequest( trace=trace, messages=[ChatMessage(role="user", content=question)], context_summary=env.payload.get("context_summary"), ) resp = await gw.chat(req) return HandlerResult( status=JobStatus.COMPLETED, result={ "kind": "ask_asf", "question": question, "answer": resp.reply, "references": resp.references, }, usage=JobUsage( input_tokens=resp.input_tokens, output_tokens=resp.output_tokens, estimated_cost=resp.estimated_cost, ), ) if kind == "analyze_code": req = CodeAnalysisRequest( trace=trace, code_snippet=env.payload.get("code_snippet", ""), language=env.payload.get("language"), file_path=env.payload.get("file_path"), ) resp = await gw.analyze_code(req) return HandlerResult( status=JobStatus.COMPLETED, result={ "kind": "analyze_code", "summary": resp.summary, "findings": [f.model_dump(mode="json") for f in resp.findings], }, usage=JobUsage( input_tokens=resp.input_tokens, output_tokens=resp.output_tokens, estimated_cost=resp.estimated_cost, ), ) if kind == "generate_tests": req = TestGenerationRequest( trace=trace, code_snippet=env.payload.get("code_snippet", ""), language=env.payload.get("language", "python"), test_framework=env.payload.get("framework", "pytest"), goal=env.payload.get("goal", "Generate unit tests"), ) resp = await gw.generate_tests(req) return HandlerResult( status=JobStatus.COMPLETED, result={ "kind": "generate_tests", "test_cases": resp.test_cases, "fixtures_needed": resp.fixtures_needed, "coverage_notes": resp.coverage_notes, }, usage=JobUsage( input_tokens=resp.input_tokens, output_tokens=resp.output_tokens, estimated_cost=resp.estimated_cost, ), ) if kind == "decompose_requirement": req = RequirementDecomposeRequest( trace=trace, requirement_text=env.payload.get("requirement_text", ""), context=env.payload.get("context"), max_subtasks=int(env.payload.get("max_subtasks", 20)), ) resp = await gw.decompose_requirement(req) return HandlerResult( status=JobStatus.COMPLETED, result={ "kind": "decompose_requirement", "subtasks": [s.model_dump(mode="json") for s in resp.subtasks], "architecture_notes": resp.architecture_notes, "open_questions": resp.open_questions, }, usage=JobUsage( input_tokens=resp.input_tokens, output_tokens=resp.output_tokens, estimated_cost=resp.estimated_cost, ), ) return await echo_handler(env, deps) except Exception as exc: # noqa: BLE001 log.exception("ai_handler_failed", exc=str(exc)) return HandlerResult( status=JobStatus.FAILED, error=f"AI handler raised: {type(exc).__name__}: {exc}", error_detail={"payload": env.payload}, ) async def indexing_handler( env: JobRequestEnvelope, deps: WorkerDependencies ) -> HandlerResult: """Stub handler for repository / doc indexing jobs.""" log.info( "indexing_handler_stub", job_id=str(env.job_id), payload=env.payload, ) return HandlerResult( status=JobStatus.COMPLETED, result={"stub": True, "indexed_items": 0, "source": env.payload.get("source")}, usage=JobUsage(cpu_seconds=0.005), ) async def notification_handler( env: JobRequestEnvelope, deps: WorkerDependencies ) -> HandlerResult: """Stub handler for email / webhook / in-app notifications.""" log.info( "notification_stub", job_id=str(env.job_id), payload_keys=list(env.payload.keys()), ) return HandlerResult( status=JobStatus.COMPLETED, result={"note": "notifications are a stub in bootstrap"}, ) # Final status persistence step, exported for reuse in the main worker loop. mark_job_done = _mark_job_done