kubiya-control-plane-api 0.9.15__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (479) hide show
  1. control_plane_api/LICENSE +676 -0
  2. control_plane_api/README.md +350 -0
  3. control_plane_api/__init__.py +4 -0
  4. control_plane_api/__version__.py +8 -0
  5. control_plane_api/alembic/README +1 -0
  6. control_plane_api/alembic/env.py +121 -0
  7. control_plane_api/alembic/script.py.mako +28 -0
  8. control_plane_api/alembic/versions/2613c65c3dbe_initial_database_setup.py +32 -0
  9. control_plane_api/alembic/versions/2df520d4927d_merge_heads.py +28 -0
  10. control_plane_api/alembic/versions/43abf98d6a01_add_paused_status_to_executions.py +73 -0
  11. control_plane_api/alembic/versions/6289854264cb_merge_multiple_heads.py +28 -0
  12. control_plane_api/alembic/versions/6a4d4dc3d8dc_generate_execution_transitions.py +50 -0
  13. control_plane_api/alembic/versions/87d11cf0a783_add_disconnected_status_to_worker_.py +44 -0
  14. control_plane_api/alembic/versions/add_ephemeral_queue_support.py +85 -0
  15. control_plane_api/alembic/versions/add_model_type_to_llm_models.py +31 -0
  16. control_plane_api/alembic/versions/add_plan_executions_table.py +114 -0
  17. control_plane_api/alembic/versions/add_trace_span_tables.py +154 -0
  18. control_plane_api/alembic/versions/add_user_info_to_traces.py +36 -0
  19. control_plane_api/alembic/versions/adjusting_foreign_keys.py +32 -0
  20. control_plane_api/alembic/versions/b4983d976db2_initial_tables.py +1128 -0
  21. control_plane_api/alembic/versions/d181a3b40e71_rename_custom_metadata_to_metadata_in_.py +50 -0
  22. control_plane_api/alembic/versions/df9117888e82_add_missing_columns.py +82 -0
  23. control_plane_api/alembic/versions/f25de6ad895a_missing_migrations.py +34 -0
  24. control_plane_api/alembic/versions/f71305fb69b9_fix_ephemeral_queue_deletion_foreign_key.py +54 -0
  25. control_plane_api/alembic/versions/mark_local_exec_queues_as_ephemeral.py +68 -0
  26. control_plane_api/alembic.ini +148 -0
  27. control_plane_api/api/index.py +12 -0
  28. control_plane_api/app/__init__.py +11 -0
  29. control_plane_api/app/activities/__init__.py +20 -0
  30. control_plane_api/app/activities/agent_activities.py +384 -0
  31. control_plane_api/app/activities/plan_generation_activities.py +499 -0
  32. control_plane_api/app/activities/team_activities.py +424 -0
  33. control_plane_api/app/activities/temporal_cloud_activities.py +588 -0
  34. control_plane_api/app/config/__init__.py +35 -0
  35. control_plane_api/app/config/api_config.py +469 -0
  36. control_plane_api/app/config/config_loader.py +224 -0
  37. control_plane_api/app/config/model_pricing.py +323 -0
  38. control_plane_api/app/config/storage_config.py +159 -0
  39. control_plane_api/app/config.py +115 -0
  40. control_plane_api/app/controllers/__init__.py +0 -0
  41. control_plane_api/app/controllers/execution_environment_controller.py +1315 -0
  42. control_plane_api/app/database.py +135 -0
  43. control_plane_api/app/exceptions.py +408 -0
  44. control_plane_api/app/lib/__init__.py +11 -0
  45. control_plane_api/app/lib/environment.py +65 -0
  46. control_plane_api/app/lib/event_bus/__init__.py +17 -0
  47. control_plane_api/app/lib/event_bus/base.py +136 -0
  48. control_plane_api/app/lib/event_bus/manager.py +335 -0
  49. control_plane_api/app/lib/event_bus/providers/__init__.py +6 -0
  50. control_plane_api/app/lib/event_bus/providers/http_provider.py +166 -0
  51. control_plane_api/app/lib/event_bus/providers/nats_provider.py +324 -0
  52. control_plane_api/app/lib/event_bus/providers/redis_provider.py +233 -0
  53. control_plane_api/app/lib/event_bus/providers/websocket_provider.py +497 -0
  54. control_plane_api/app/lib/job_executor.py +330 -0
  55. control_plane_api/app/lib/kubiya_client.py +293 -0
  56. control_plane_api/app/lib/litellm_pricing.py +166 -0
  57. control_plane_api/app/lib/mcp_validation.py +163 -0
  58. control_plane_api/app/lib/nats/__init__.py +13 -0
  59. control_plane_api/app/lib/nats/credentials_manager.py +288 -0
  60. control_plane_api/app/lib/nats/listener.py +374 -0
  61. control_plane_api/app/lib/planning_prompt_builder.py +153 -0
  62. control_plane_api/app/lib/planning_tools/__init__.py +41 -0
  63. control_plane_api/app/lib/planning_tools/agents.py +409 -0
  64. control_plane_api/app/lib/planning_tools/agno_toolkit.py +836 -0
  65. control_plane_api/app/lib/planning_tools/base.py +119 -0
  66. control_plane_api/app/lib/planning_tools/cognitive_memory_tools.py +403 -0
  67. control_plane_api/app/lib/planning_tools/context_graph_tools.py +545 -0
  68. control_plane_api/app/lib/planning_tools/environments.py +218 -0
  69. control_plane_api/app/lib/planning_tools/knowledge.py +204 -0
  70. control_plane_api/app/lib/planning_tools/models.py +93 -0
  71. control_plane_api/app/lib/planning_tools/planning_service.py +646 -0
  72. control_plane_api/app/lib/planning_tools/resources.py +242 -0
  73. control_plane_api/app/lib/planning_tools/teams.py +334 -0
  74. control_plane_api/app/lib/policy_enforcer_client.py +1016 -0
  75. control_plane_api/app/lib/redis_client.py +803 -0
  76. control_plane_api/app/lib/sqlalchemy_utils.py +486 -0
  77. control_plane_api/app/lib/state_transition_tools/__init__.py +7 -0
  78. control_plane_api/app/lib/state_transition_tools/execution_context.py +388 -0
  79. control_plane_api/app/lib/storage/__init__.py +20 -0
  80. control_plane_api/app/lib/storage/base_provider.py +274 -0
  81. control_plane_api/app/lib/storage/provider_factory.py +157 -0
  82. control_plane_api/app/lib/storage/vercel_blob_provider.py +468 -0
  83. control_plane_api/app/lib/supabase.py +71 -0
  84. control_plane_api/app/lib/supabase_utils.py +138 -0
  85. control_plane_api/app/lib/task_planning/__init__.py +138 -0
  86. control_plane_api/app/lib/task_planning/agent_factory.py +308 -0
  87. control_plane_api/app/lib/task_planning/agents.py +389 -0
  88. control_plane_api/app/lib/task_planning/cache.py +218 -0
  89. control_plane_api/app/lib/task_planning/entity_resolver.py +273 -0
  90. control_plane_api/app/lib/task_planning/helpers.py +293 -0
  91. control_plane_api/app/lib/task_planning/hooks.py +474 -0
  92. control_plane_api/app/lib/task_planning/models.py +503 -0
  93. control_plane_api/app/lib/task_planning/plan_validator.py +166 -0
  94. control_plane_api/app/lib/task_planning/planning_workflow.py +2911 -0
  95. control_plane_api/app/lib/task_planning/runner.py +656 -0
  96. control_plane_api/app/lib/task_planning/streaming_hook.py +213 -0
  97. control_plane_api/app/lib/task_planning/workflow.py +424 -0
  98. control_plane_api/app/lib/templating/__init__.py +88 -0
  99. control_plane_api/app/lib/templating/compiler.py +278 -0
  100. control_plane_api/app/lib/templating/engine.py +178 -0
  101. control_plane_api/app/lib/templating/parsers/__init__.py +29 -0
  102. control_plane_api/app/lib/templating/parsers/base.py +96 -0
  103. control_plane_api/app/lib/templating/parsers/env.py +85 -0
  104. control_plane_api/app/lib/templating/parsers/graph.py +112 -0
  105. control_plane_api/app/lib/templating/parsers/secret.py +87 -0
  106. control_plane_api/app/lib/templating/parsers/simple.py +81 -0
  107. control_plane_api/app/lib/templating/resolver.py +366 -0
  108. control_plane_api/app/lib/templating/types.py +214 -0
  109. control_plane_api/app/lib/templating/validator.py +201 -0
  110. control_plane_api/app/lib/temporal_client.py +232 -0
  111. control_plane_api/app/lib/temporal_credentials_cache.py +178 -0
  112. control_plane_api/app/lib/temporal_credentials_service.py +203 -0
  113. control_plane_api/app/lib/validation/__init__.py +24 -0
  114. control_plane_api/app/lib/validation/runtime_validation.py +388 -0
  115. control_plane_api/app/main.py +531 -0
  116. control_plane_api/app/middleware/__init__.py +10 -0
  117. control_plane_api/app/middleware/auth.py +645 -0
  118. control_plane_api/app/middleware/exception_handler.py +267 -0
  119. control_plane_api/app/middleware/prometheus_middleware.py +173 -0
  120. control_plane_api/app/middleware/rate_limiting.py +384 -0
  121. control_plane_api/app/middleware/request_id.py +202 -0
  122. control_plane_api/app/models/__init__.py +40 -0
  123. control_plane_api/app/models/agent.py +90 -0
  124. control_plane_api/app/models/analytics.py +206 -0
  125. control_plane_api/app/models/associations.py +107 -0
  126. control_plane_api/app/models/auth_user.py +73 -0
  127. control_plane_api/app/models/context.py +161 -0
  128. control_plane_api/app/models/custom_integration.py +99 -0
  129. control_plane_api/app/models/environment.py +64 -0
  130. control_plane_api/app/models/execution.py +125 -0
  131. control_plane_api/app/models/execution_transition.py +50 -0
  132. control_plane_api/app/models/job.py +159 -0
  133. control_plane_api/app/models/llm_model.py +78 -0
  134. control_plane_api/app/models/orchestration.py +66 -0
  135. control_plane_api/app/models/plan_execution.py +102 -0
  136. control_plane_api/app/models/presence.py +49 -0
  137. control_plane_api/app/models/project.py +61 -0
  138. control_plane_api/app/models/project_management.py +85 -0
  139. control_plane_api/app/models/session.py +29 -0
  140. control_plane_api/app/models/skill.py +155 -0
  141. control_plane_api/app/models/system_tables.py +43 -0
  142. control_plane_api/app/models/task_planning.py +372 -0
  143. control_plane_api/app/models/team.py +86 -0
  144. control_plane_api/app/models/trace.py +257 -0
  145. control_plane_api/app/models/user_profile.py +54 -0
  146. control_plane_api/app/models/worker.py +221 -0
  147. control_plane_api/app/models/workflow.py +161 -0
  148. control_plane_api/app/models/workspace.py +50 -0
  149. control_plane_api/app/observability/__init__.py +177 -0
  150. control_plane_api/app/observability/context_logging.py +475 -0
  151. control_plane_api/app/observability/decorators.py +337 -0
  152. control_plane_api/app/observability/local_span_processor.py +702 -0
  153. control_plane_api/app/observability/metrics.py +303 -0
  154. control_plane_api/app/observability/middleware.py +246 -0
  155. control_plane_api/app/observability/optional.py +115 -0
  156. control_plane_api/app/observability/tracing.py +382 -0
  157. control_plane_api/app/policies/README.md +149 -0
  158. control_plane_api/app/policies/approved_users.rego +62 -0
  159. control_plane_api/app/policies/business_hours.rego +51 -0
  160. control_plane_api/app/policies/rate_limiting.rego +100 -0
  161. control_plane_api/app/policies/tool_enforcement/README.md +336 -0
  162. control_plane_api/app/policies/tool_enforcement/bash_command_validation.rego +71 -0
  163. control_plane_api/app/policies/tool_enforcement/business_hours_enforcement.rego +82 -0
  164. control_plane_api/app/policies/tool_enforcement/mcp_tool_allowlist.rego +58 -0
  165. control_plane_api/app/policies/tool_enforcement/production_safeguards.rego +80 -0
  166. control_plane_api/app/policies/tool_enforcement/role_based_tool_access.rego +44 -0
  167. control_plane_api/app/policies/tool_restrictions.rego +86 -0
  168. control_plane_api/app/routers/__init__.py +4 -0
  169. control_plane_api/app/routers/agents.py +382 -0
  170. control_plane_api/app/routers/agents_v2.py +1598 -0
  171. control_plane_api/app/routers/analytics.py +1310 -0
  172. control_plane_api/app/routers/auth.py +59 -0
  173. control_plane_api/app/routers/client_config.py +57 -0
  174. control_plane_api/app/routers/context_graph.py +561 -0
  175. control_plane_api/app/routers/context_manager.py +577 -0
  176. control_plane_api/app/routers/custom_integrations.py +490 -0
  177. control_plane_api/app/routers/enforcer.py +132 -0
  178. control_plane_api/app/routers/environment_context.py +252 -0
  179. control_plane_api/app/routers/environments.py +761 -0
  180. control_plane_api/app/routers/execution_environment.py +847 -0
  181. control_plane_api/app/routers/executions/__init__.py +28 -0
  182. control_plane_api/app/routers/executions/router.py +286 -0
  183. control_plane_api/app/routers/executions/services/__init__.py +22 -0
  184. control_plane_api/app/routers/executions/services/demo_worker_health.py +156 -0
  185. control_plane_api/app/routers/executions/services/status_service.py +420 -0
  186. control_plane_api/app/routers/executions/services/test_worker_health.py +480 -0
  187. control_plane_api/app/routers/executions/services/worker_health.py +514 -0
  188. control_plane_api/app/routers/executions/streaming/__init__.py +22 -0
  189. control_plane_api/app/routers/executions/streaming/deduplication.py +352 -0
  190. control_plane_api/app/routers/executions/streaming/event_buffer.py +353 -0
  191. control_plane_api/app/routers/executions/streaming/event_formatter.py +964 -0
  192. control_plane_api/app/routers/executions/streaming/history_loader.py +588 -0
  193. control_plane_api/app/routers/executions/streaming/live_source.py +693 -0
  194. control_plane_api/app/routers/executions/streaming/streamer.py +849 -0
  195. control_plane_api/app/routers/executions.py +4888 -0
  196. control_plane_api/app/routers/health.py +165 -0
  197. control_plane_api/app/routers/health_v2.py +394 -0
  198. control_plane_api/app/routers/integration_templates.py +496 -0
  199. control_plane_api/app/routers/integrations.py +287 -0
  200. control_plane_api/app/routers/jobs.py +1809 -0
  201. control_plane_api/app/routers/metrics.py +517 -0
  202. control_plane_api/app/routers/models.py +82 -0
  203. control_plane_api/app/routers/models_v2.py +628 -0
  204. control_plane_api/app/routers/plan_executions.py +1481 -0
  205. control_plane_api/app/routers/plan_generation_async.py +304 -0
  206. control_plane_api/app/routers/policies.py +669 -0
  207. control_plane_api/app/routers/presence.py +234 -0
  208. control_plane_api/app/routers/projects.py +987 -0
  209. control_plane_api/app/routers/runners.py +379 -0
  210. control_plane_api/app/routers/runtimes.py +172 -0
  211. control_plane_api/app/routers/secrets.py +171 -0
  212. control_plane_api/app/routers/skills.py +1010 -0
  213. control_plane_api/app/routers/skills_definitions.py +140 -0
  214. control_plane_api/app/routers/storage.py +456 -0
  215. control_plane_api/app/routers/task_planning.py +611 -0
  216. control_plane_api/app/routers/task_queues.py +650 -0
  217. control_plane_api/app/routers/team_context.py +274 -0
  218. control_plane_api/app/routers/teams.py +1747 -0
  219. control_plane_api/app/routers/templates.py +248 -0
  220. control_plane_api/app/routers/traces.py +571 -0
  221. control_plane_api/app/routers/websocket_client.py +479 -0
  222. control_plane_api/app/routers/websocket_executions_status.py +437 -0
  223. control_plane_api/app/routers/websocket_gateway.py +323 -0
  224. control_plane_api/app/routers/websocket_traces.py +576 -0
  225. control_plane_api/app/routers/worker_queues.py +2555 -0
  226. control_plane_api/app/routers/worker_websocket.py +419 -0
  227. control_plane_api/app/routers/workers.py +1004 -0
  228. control_plane_api/app/routers/workflows.py +204 -0
  229. control_plane_api/app/runtimes/__init__.py +6 -0
  230. control_plane_api/app/runtimes/validation.py +344 -0
  231. control_plane_api/app/schemas/__init__.py +1 -0
  232. control_plane_api/app/schemas/job_schemas.py +302 -0
  233. control_plane_api/app/schemas/mcp_schemas.py +311 -0
  234. control_plane_api/app/schemas/template_schemas.py +133 -0
  235. control_plane_api/app/schemas/trace_schemas.py +168 -0
  236. control_plane_api/app/schemas/worker_queue_observability_schemas.py +165 -0
  237. control_plane_api/app/services/__init__.py +1 -0
  238. control_plane_api/app/services/agno_planning_strategy.py +233 -0
  239. control_plane_api/app/services/agno_service.py +838 -0
  240. control_plane_api/app/services/claude_code_planning_service.py +203 -0
  241. control_plane_api/app/services/context_graph_client.py +224 -0
  242. control_plane_api/app/services/custom_integration_service.py +415 -0
  243. control_plane_api/app/services/integration_resolution_service.py +345 -0
  244. control_plane_api/app/services/litellm_service.py +394 -0
  245. control_plane_api/app/services/plan_generator.py +79 -0
  246. control_plane_api/app/services/planning_strategy.py +66 -0
  247. control_plane_api/app/services/planning_strategy_factory.py +118 -0
  248. control_plane_api/app/services/policy_service.py +615 -0
  249. control_plane_api/app/services/state_transition_service.py +755 -0
  250. control_plane_api/app/services/storage_service.py +593 -0
  251. control_plane_api/app/services/temporal_cloud_provisioning.py +150 -0
  252. control_plane_api/app/services/toolsets/context_graph_skill.py +432 -0
  253. control_plane_api/app/services/trace_retention.py +354 -0
  254. control_plane_api/app/services/worker_queue_metrics_service.py +190 -0
  255. control_plane_api/app/services/workflow_cancellation_manager.py +135 -0
  256. control_plane_api/app/services/workflow_operations_service.py +611 -0
  257. control_plane_api/app/skills/__init__.py +100 -0
  258. control_plane_api/app/skills/base.py +239 -0
  259. control_plane_api/app/skills/builtin/__init__.py +37 -0
  260. control_plane_api/app/skills/builtin/agent_communication/__init__.py +8 -0
  261. control_plane_api/app/skills/builtin/agent_communication/skill.py +246 -0
  262. control_plane_api/app/skills/builtin/code_ingestion/__init__.py +4 -0
  263. control_plane_api/app/skills/builtin/code_ingestion/skill.py +267 -0
  264. control_plane_api/app/skills/builtin/cognitive_memory/__init__.py +4 -0
  265. control_plane_api/app/skills/builtin/cognitive_memory/skill.py +174 -0
  266. control_plane_api/app/skills/builtin/contextual_awareness/__init__.py +4 -0
  267. control_plane_api/app/skills/builtin/contextual_awareness/skill.py +387 -0
  268. control_plane_api/app/skills/builtin/data_visualization/__init__.py +4 -0
  269. control_plane_api/app/skills/builtin/data_visualization/skill.py +154 -0
  270. control_plane_api/app/skills/builtin/docker/__init__.py +4 -0
  271. control_plane_api/app/skills/builtin/docker/skill.py +104 -0
  272. control_plane_api/app/skills/builtin/file_generation/__init__.py +4 -0
  273. control_plane_api/app/skills/builtin/file_generation/skill.py +94 -0
  274. control_plane_api/app/skills/builtin/file_system/__init__.py +4 -0
  275. control_plane_api/app/skills/builtin/file_system/skill.py +110 -0
  276. control_plane_api/app/skills/builtin/knowledge_api/__init__.py +5 -0
  277. control_plane_api/app/skills/builtin/knowledge_api/skill.py +124 -0
  278. control_plane_api/app/skills/builtin/python/__init__.py +4 -0
  279. control_plane_api/app/skills/builtin/python/skill.py +92 -0
  280. control_plane_api/app/skills/builtin/remote_filesystem/__init__.py +5 -0
  281. control_plane_api/app/skills/builtin/remote_filesystem/skill.py +170 -0
  282. control_plane_api/app/skills/builtin/shell/__init__.py +4 -0
  283. control_plane_api/app/skills/builtin/shell/skill.py +161 -0
  284. control_plane_api/app/skills/builtin/slack/__init__.py +3 -0
  285. control_plane_api/app/skills/builtin/slack/skill.py +302 -0
  286. control_plane_api/app/skills/builtin/workflow_executor/__init__.py +4 -0
  287. control_plane_api/app/skills/builtin/workflow_executor/skill.py +469 -0
  288. control_plane_api/app/skills/business_intelligence.py +189 -0
  289. control_plane_api/app/skills/config.py +63 -0
  290. control_plane_api/app/skills/loaders/__init__.py +14 -0
  291. control_plane_api/app/skills/loaders/base.py +73 -0
  292. control_plane_api/app/skills/loaders/filesystem_loader.py +199 -0
  293. control_plane_api/app/skills/registry.py +125 -0
  294. control_plane_api/app/utils/helpers.py +12 -0
  295. control_plane_api/app/utils/workflow_executor.py +354 -0
  296. control_plane_api/app/workflows/__init__.py +11 -0
  297. control_plane_api/app/workflows/agent_execution.py +520 -0
  298. control_plane_api/app/workflows/agent_execution_with_skills.py +223 -0
  299. control_plane_api/app/workflows/namespace_provisioning.py +326 -0
  300. control_plane_api/app/workflows/plan_generation.py +254 -0
  301. control_plane_api/app/workflows/team_execution.py +442 -0
  302. control_plane_api/scripts/seed_models.py +240 -0
  303. control_plane_api/scripts/validate_existing_tool_names.py +492 -0
  304. control_plane_api/shared/__init__.py +8 -0
  305. control_plane_api/shared/version.py +17 -0
  306. control_plane_api/test_deduplication.py +274 -0
  307. control_plane_api/test_executor_deduplication_e2e.py +309 -0
  308. control_plane_api/test_job_execution_e2e.py +283 -0
  309. control_plane_api/test_real_integration.py +193 -0
  310. control_plane_api/version.py +38 -0
  311. control_plane_api/worker/__init__.py +0 -0
  312. control_plane_api/worker/activities/__init__.py +0 -0
  313. control_plane_api/worker/activities/agent_activities.py +1585 -0
  314. control_plane_api/worker/activities/approval_activities.py +234 -0
  315. control_plane_api/worker/activities/job_activities.py +199 -0
  316. control_plane_api/worker/activities/runtime_activities.py +1167 -0
  317. control_plane_api/worker/activities/skill_activities.py +282 -0
  318. control_plane_api/worker/activities/team_activities.py +479 -0
  319. control_plane_api/worker/agent_runtime_server.py +370 -0
  320. control_plane_api/worker/binary_manager.py +333 -0
  321. control_plane_api/worker/config/__init__.py +31 -0
  322. control_plane_api/worker/config/worker_config.py +273 -0
  323. control_plane_api/worker/control_plane_client.py +1491 -0
  324. control_plane_api/worker/examples/analytics_integration_example.py +362 -0
  325. control_plane_api/worker/health_monitor.py +159 -0
  326. control_plane_api/worker/metrics.py +237 -0
  327. control_plane_api/worker/models/__init__.py +1 -0
  328. control_plane_api/worker/models/error_events.py +105 -0
  329. control_plane_api/worker/models/inputs.py +89 -0
  330. control_plane_api/worker/runtimes/__init__.py +35 -0
  331. control_plane_api/worker/runtimes/agent_runtime/runtime.py +485 -0
  332. control_plane_api/worker/runtimes/agno/__init__.py +34 -0
  333. control_plane_api/worker/runtimes/agno/config.py +248 -0
  334. control_plane_api/worker/runtimes/agno/hooks.py +385 -0
  335. control_plane_api/worker/runtimes/agno/mcp_builder.py +195 -0
  336. control_plane_api/worker/runtimes/agno/runtime.py +1063 -0
  337. control_plane_api/worker/runtimes/agno/utils.py +163 -0
  338. control_plane_api/worker/runtimes/base.py +979 -0
  339. control_plane_api/worker/runtimes/claude_code/__init__.py +38 -0
  340. control_plane_api/worker/runtimes/claude_code/cleanup.py +184 -0
  341. control_plane_api/worker/runtimes/claude_code/client_pool.py +529 -0
  342. control_plane_api/worker/runtimes/claude_code/config.py +829 -0
  343. control_plane_api/worker/runtimes/claude_code/hooks.py +482 -0
  344. control_plane_api/worker/runtimes/claude_code/litellm_proxy.py +1702 -0
  345. control_plane_api/worker/runtimes/claude_code/mcp_builder.py +467 -0
  346. control_plane_api/worker/runtimes/claude_code/mcp_discovery.py +558 -0
  347. control_plane_api/worker/runtimes/claude_code/runtime.py +1546 -0
  348. control_plane_api/worker/runtimes/claude_code/tool_mapper.py +403 -0
  349. control_plane_api/worker/runtimes/claude_code/utils.py +149 -0
  350. control_plane_api/worker/runtimes/factory.py +173 -0
  351. control_plane_api/worker/runtimes/model_utils.py +107 -0
  352. control_plane_api/worker/runtimes/validation.py +93 -0
  353. control_plane_api/worker/services/__init__.py +1 -0
  354. control_plane_api/worker/services/agent_communication_tools.py +908 -0
  355. control_plane_api/worker/services/agent_executor.py +485 -0
  356. control_plane_api/worker/services/agent_executor_v2.py +793 -0
  357. control_plane_api/worker/services/analytics_collector.py +457 -0
  358. control_plane_api/worker/services/analytics_service.py +464 -0
  359. control_plane_api/worker/services/approval_tools.py +310 -0
  360. control_plane_api/worker/services/approval_tools_agno.py +207 -0
  361. control_plane_api/worker/services/cancellation_manager.py +177 -0
  362. control_plane_api/worker/services/code_ingestion_tools.py +465 -0
  363. control_plane_api/worker/services/contextual_awareness_tools.py +405 -0
  364. control_plane_api/worker/services/data_visualization.py +834 -0
  365. control_plane_api/worker/services/event_publisher.py +531 -0
  366. control_plane_api/worker/services/jira_tools.py +257 -0
  367. control_plane_api/worker/services/remote_filesystem_tools.py +498 -0
  368. control_plane_api/worker/services/runtime_analytics.py +328 -0
  369. control_plane_api/worker/services/session_service.py +365 -0
  370. control_plane_api/worker/services/skill_context_enhancement.py +181 -0
  371. control_plane_api/worker/services/skill_factory.py +471 -0
  372. control_plane_api/worker/services/system_prompt_enhancement.py +410 -0
  373. control_plane_api/worker/services/team_executor.py +715 -0
  374. control_plane_api/worker/services/team_executor_v2.py +1866 -0
  375. control_plane_api/worker/services/tool_enforcement.py +254 -0
  376. control_plane_api/worker/services/workflow_executor/__init__.py +52 -0
  377. control_plane_api/worker/services/workflow_executor/event_processor.py +287 -0
  378. control_plane_api/worker/services/workflow_executor/event_publisher.py +210 -0
  379. control_plane_api/worker/services/workflow_executor/executors/__init__.py +15 -0
  380. control_plane_api/worker/services/workflow_executor/executors/base.py +270 -0
  381. control_plane_api/worker/services/workflow_executor/executors/json_executor.py +50 -0
  382. control_plane_api/worker/services/workflow_executor/executors/python_executor.py +50 -0
  383. control_plane_api/worker/services/workflow_executor/models.py +142 -0
  384. control_plane_api/worker/services/workflow_executor_tools.py +1748 -0
  385. control_plane_api/worker/skills/__init__.py +12 -0
  386. control_plane_api/worker/skills/builtin/context_graph_search/README.md +213 -0
  387. control_plane_api/worker/skills/builtin/context_graph_search/__init__.py +5 -0
  388. control_plane_api/worker/skills/builtin/context_graph_search/agno_impl.py +808 -0
  389. control_plane_api/worker/skills/builtin/context_graph_search/skill.yaml +67 -0
  390. control_plane_api/worker/skills/builtin/contextual_awareness/__init__.py +4 -0
  391. control_plane_api/worker/skills/builtin/contextual_awareness/agno_impl.py +62 -0
  392. control_plane_api/worker/skills/builtin/data_visualization/agno_impl.py +18 -0
  393. control_plane_api/worker/skills/builtin/data_visualization/skill.yaml +84 -0
  394. control_plane_api/worker/skills/builtin/docker/agno_impl.py +65 -0
  395. control_plane_api/worker/skills/builtin/docker/skill.yaml +60 -0
  396. control_plane_api/worker/skills/builtin/file_generation/agno_impl.py +47 -0
  397. control_plane_api/worker/skills/builtin/file_generation/skill.yaml +64 -0
  398. control_plane_api/worker/skills/builtin/file_system/agno_impl.py +32 -0
  399. control_plane_api/worker/skills/builtin/file_system/skill.yaml +54 -0
  400. control_plane_api/worker/skills/builtin/knowledge_api/__init__.py +4 -0
  401. control_plane_api/worker/skills/builtin/knowledge_api/agno_impl.py +50 -0
  402. control_plane_api/worker/skills/builtin/knowledge_api/skill.yaml +66 -0
  403. control_plane_api/worker/skills/builtin/python/agno_impl.py +25 -0
  404. control_plane_api/worker/skills/builtin/python/skill.yaml +60 -0
  405. control_plane_api/worker/skills/builtin/schema_fix_mixin.py +260 -0
  406. control_plane_api/worker/skills/builtin/shell/agno_impl.py +31 -0
  407. control_plane_api/worker/skills/builtin/shell/skill.yaml +60 -0
  408. control_plane_api/worker/skills/builtin/slack/__init__.py +3 -0
  409. control_plane_api/worker/skills/builtin/slack/agno_impl.py +1282 -0
  410. control_plane_api/worker/skills/builtin/slack/skill.yaml +276 -0
  411. control_plane_api/worker/skills/builtin/workflow_executor/agno_impl.py +62 -0
  412. control_plane_api/worker/skills/builtin/workflow_executor/skill.yaml +79 -0
  413. control_plane_api/worker/skills/loaders/__init__.py +5 -0
  414. control_plane_api/worker/skills/loaders/base.py +23 -0
  415. control_plane_api/worker/skills/loaders/filesystem_loader.py +357 -0
  416. control_plane_api/worker/skills/registry.py +208 -0
  417. control_plane_api/worker/tests/__init__.py +1 -0
  418. control_plane_api/worker/tests/conftest.py +12 -0
  419. control_plane_api/worker/tests/e2e/__init__.py +0 -0
  420. control_plane_api/worker/tests/e2e/test_context_graph_real_api.py +338 -0
  421. control_plane_api/worker/tests/e2e/test_context_graph_templates_e2e.py +523 -0
  422. control_plane_api/worker/tests/e2e/test_enforcement_e2e.py +344 -0
  423. control_plane_api/worker/tests/e2e/test_execution_flow.py +571 -0
  424. control_plane_api/worker/tests/e2e/test_single_execution_mode.py +656 -0
  425. control_plane_api/worker/tests/integration/__init__.py +0 -0
  426. control_plane_api/worker/tests/integration/test_builtin_skills_fixes.py +245 -0
  427. control_plane_api/worker/tests/integration/test_context_graph_search_integration.py +365 -0
  428. control_plane_api/worker/tests/integration/test_control_plane_integration.py +308 -0
  429. control_plane_api/worker/tests/integration/test_hook_enforcement_integration.py +579 -0
  430. control_plane_api/worker/tests/integration/test_scheduled_job_workflow.py +237 -0
  431. control_plane_api/worker/tests/integration/test_system_prompt_enhancement_integration.py +343 -0
  432. control_plane_api/worker/tests/unit/__init__.py +0 -0
  433. control_plane_api/worker/tests/unit/test_builtin_skill_autoload.py +396 -0
  434. control_plane_api/worker/tests/unit/test_context_graph_search.py +450 -0
  435. control_plane_api/worker/tests/unit/test_context_graph_templates.py +403 -0
  436. control_plane_api/worker/tests/unit/test_control_plane_client.py +401 -0
  437. control_plane_api/worker/tests/unit/test_control_plane_client_jobs.py +345 -0
  438. control_plane_api/worker/tests/unit/test_job_activities.py +353 -0
  439. control_plane_api/worker/tests/unit/test_skill_context_enhancement.py +321 -0
  440. control_plane_api/worker/tests/unit/test_system_prompt_enhancement.py +415 -0
  441. control_plane_api/worker/tests/unit/test_tool_enforcement.py +324 -0
  442. control_plane_api/worker/utils/__init__.py +1 -0
  443. control_plane_api/worker/utils/chunk_batcher.py +330 -0
  444. control_plane_api/worker/utils/environment.py +65 -0
  445. control_plane_api/worker/utils/error_publisher.py +260 -0
  446. control_plane_api/worker/utils/event_batcher.py +256 -0
  447. control_plane_api/worker/utils/logging_config.py +335 -0
  448. control_plane_api/worker/utils/logging_helper.py +326 -0
  449. control_plane_api/worker/utils/parameter_validator.py +120 -0
  450. control_plane_api/worker/utils/retry_utils.py +60 -0
  451. control_plane_api/worker/utils/streaming_utils.py +665 -0
  452. control_plane_api/worker/utils/tool_validation.py +332 -0
  453. control_plane_api/worker/utils/workspace_manager.py +163 -0
  454. control_plane_api/worker/websocket_client.py +393 -0
  455. control_plane_api/worker/worker.py +1297 -0
  456. control_plane_api/worker/workflows/__init__.py +0 -0
  457. control_plane_api/worker/workflows/agent_execution.py +909 -0
  458. control_plane_api/worker/workflows/scheduled_job_wrapper.py +332 -0
  459. control_plane_api/worker/workflows/team_execution.py +611 -0
  460. kubiya_control_plane_api-0.9.15.dist-info/METADATA +354 -0
  461. kubiya_control_plane_api-0.9.15.dist-info/RECORD +479 -0
  462. kubiya_control_plane_api-0.9.15.dist-info/WHEEL +5 -0
  463. kubiya_control_plane_api-0.9.15.dist-info/entry_points.txt +5 -0
  464. kubiya_control_plane_api-0.9.15.dist-info/licenses/LICENSE +676 -0
  465. kubiya_control_plane_api-0.9.15.dist-info/top_level.txt +3 -0
  466. scripts/__init__.py +1 -0
  467. scripts/migrations.py +39 -0
  468. scripts/seed_worker_queues.py +128 -0
  469. scripts/setup_agent_runtime.py +142 -0
  470. worker_internal/__init__.py +1 -0
  471. worker_internal/planner/__init__.py +1 -0
  472. worker_internal/planner/activities.py +1499 -0
  473. worker_internal/planner/agent_tools.py +197 -0
  474. worker_internal/planner/event_models.py +148 -0
  475. worker_internal/planner/event_publisher.py +67 -0
  476. worker_internal/planner/models.py +199 -0
  477. worker_internal/planner/retry_logic.py +134 -0
  478. worker_internal/planner/worker.py +300 -0
  479. worker_internal/planner/workflows.py +970 -0
@@ -0,0 +1,1016 @@
1
+ """
2
+ Policy Enforcer Client - Integration with OPA Watchdog Enforcer Service.
3
+
4
+ This module provides a robust, async client for interacting with the OPA Watchdog
5
+ policy enforcement service. It follows best practices including:
6
+
7
+ - Async/await for non-blocking I/O
8
+ - Proper exception hierarchy
9
+ - Context manager support for resource cleanup
10
+ - Pydantic models for validation
11
+ - Dependency injection (no singletons)
12
+ - Separation of concerns with specialized clients
13
+ - Retry logic and circuit breaker patterns
14
+ - Comprehensive logging
15
+
16
+ The enforcer service URL is configured via ENFORCER_SERVICE_URL environment variable.
17
+ Default: https://enforcer-psi.vercel.app
18
+
19
+ Usage:
20
+ from control_plane_api.app.lib.policy_enforcer_client import PolicyEnforcerClient
21
+
22
+ async with PolicyEnforcerClient(base_url="...", api_key="...") as client:
23
+ policy = await client.policies.create(name="...", policy_content="...")
24
+ result = await client.evaluation.evaluate(policy_id="...", input_data={...})
25
+ """
26
+
27
+ import os
28
+ import httpx
29
+ from datetime import datetime, timezone
30
+ from typing import Optional, Dict, List, Any, Protocol
31
+ from enum import Enum
32
+ from pydantic import BaseModel, Field, validator
33
+ import structlog
34
+ from contextlib import asynccontextmanager
35
+
36
+ logger = structlog.get_logger()
37
+
38
+
39
+ # ============================================================================
40
+ # Custom Exceptions
41
+ # ============================================================================
42
+
43
+ class PolicyEnforcerError(Exception):
44
+ """Base exception for all policy enforcer errors"""
45
+ pass
46
+
47
+
48
+ class PolicyNotFoundError(PolicyEnforcerError):
49
+ """Raised when a policy is not found"""
50
+ pass
51
+
52
+
53
+ class PolicyValidationError(PolicyEnforcerError):
54
+ """Raised when policy validation fails"""
55
+ def __init__(self, message: str, details: dict = None, errors: list = None):
56
+ super().__init__(message)
57
+ self.details = details or {}
58
+ self.errors = errors or []
59
+
60
+
61
+ class PolicyEvaluationError(PolicyEnforcerError):
62
+ """Raised when policy evaluation fails"""
63
+ pass
64
+
65
+
66
+ class RequestNotFoundError(PolicyEnforcerError):
67
+ """Raised when a request is not found"""
68
+ pass
69
+
70
+
71
+ class EnforcerConnectionError(PolicyEnforcerError):
72
+ """Raised when connection to enforcer service fails"""
73
+ pass
74
+
75
+
76
+ class EnforcerAuthenticationError(PolicyEnforcerError):
77
+ """Raised when authentication with enforcer service fails"""
78
+ pass
79
+
80
+
81
+ # ============================================================================
82
+ # Pydantic Models
83
+ # ============================================================================
84
+
85
+ class PolicyType(str, Enum):
86
+ """Policy type enumeration"""
87
+ REGO = "rego"
88
+
89
+
90
+ class Decision(str, Enum):
91
+ """Policy evaluation decision"""
92
+ PERMIT = "permit"
93
+ DENY = "deny"
94
+
95
+
96
+ class RequestStatus(str, Enum):
97
+ """Request approval status"""
98
+ PENDING = "pending"
99
+ APPROVED = "approved"
100
+ REJECTED = "rejected"
101
+ EXPIRED = "expired"
102
+
103
+
104
+ class Policy(BaseModel):
105
+ """Policy model matching the enforcer service schema"""
106
+ id: str = Field(..., description="Policy UUID")
107
+ name: str = Field(..., min_length=1, max_length=255, description="Policy name")
108
+ policy_content: Optional[str] = Field(None, description="OPA Rego policy content (optional in list responses)")
109
+ org: str = Field(..., description="Organization ID")
110
+ enabled: bool = Field(default=True, description="Whether policy is enabled")
111
+ description: Optional[str] = Field(None, description="Policy description")
112
+ policy_type: PolicyType = Field(default=PolicyType.REGO, description="Policy type")
113
+ tags: List[str] = Field(default_factory=list, description="Policy tags")
114
+ metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
115
+ version: int = Field(default=1, ge=1, description="Policy version")
116
+ created_at: Optional[datetime] = Field(None, description="Creation timestamp")
117
+ updated_at: Optional[datetime] = Field(None, description="Last update timestamp")
118
+ created_by: Optional[str] = Field(None, description="Creator email")
119
+ updated_by: Optional[str] = Field(None, description="Last updater email")
120
+
121
+ class Config:
122
+ use_enum_values = True
123
+
124
+
125
+ class PolicyCreate(BaseModel):
126
+ """Schema for creating a new policy"""
127
+ name: str = Field(..., min_length=1, max_length=255)
128
+ policy_content: str = Field(..., min_length=1)
129
+ description: Optional[str] = None
130
+ enabled: bool = True
131
+ tags: List[str] = Field(default_factory=list)
132
+ metadata: Dict[str, Any] = Field(default_factory=dict)
133
+
134
+
135
+ class PolicyUpdate(BaseModel):
136
+ """Schema for updating an existing policy"""
137
+ name: Optional[str] = Field(None, min_length=1, max_length=255)
138
+ policy_content: Optional[str] = None
139
+ description: Optional[str] = None
140
+ enabled: Optional[bool] = None
141
+ tags: Optional[List[str]] = None
142
+ metadata: Optional[Dict[str, Any]] = None
143
+
144
+
145
+ class EvaluationResult(BaseModel):
146
+ """Policy evaluation result"""
147
+ allow: bool = Field(..., description="Whether the action is allowed")
148
+ decision: Decision = Field(..., description="Evaluation decision")
149
+ violations: List[str] = Field(default_factory=list, description="List of violations")
150
+ metadata: Dict[str, Any] = Field(default_factory=dict, description="Evaluation metadata")
151
+
152
+ class Config:
153
+ use_enum_values = True
154
+
155
+
156
+ class EvaluationRequest(BaseModel):
157
+ """Policy evaluation request"""
158
+ input: Dict[str, Any] = Field(..., description="Input data for evaluation")
159
+ policy_id: Optional[str] = Field(None, description="Policy UUID")
160
+ policy_name: Optional[str] = Field(None, description="Policy name")
161
+
162
+ @validator("policy_id", "policy_name")
163
+ def validate_policy_identifier(cls, v, values):
164
+ """Ensure at least one policy identifier is provided"""
165
+ if not v and not values.get("policy_id") and not values.get("policy_name"):
166
+ raise ValueError("Either policy_id or policy_name must be provided")
167
+ return v
168
+
169
+
170
+ class ValidationResult(BaseModel):
171
+ """Policy validation result"""
172
+ valid: bool = Field(..., description="Whether the policy is valid")
173
+ errors: List[str] = Field(default_factory=list, description="Validation errors")
174
+ warnings: List[str] = Field(default_factory=list, description="Validation warnings")
175
+
176
+
177
+ class ApprovalRequest(BaseModel):
178
+ """Approval request model"""
179
+ id: str = Field(..., description="Request UUID")
180
+ request_id: str = Field(..., description="Request identifier")
181
+ org: str = Field(..., description="Organization ID")
182
+ runner: str = Field(..., description="Runner name")
183
+ request_hash: str = Field(..., description="Request hash")
184
+ approved: bool = Field(..., description="Whether approved")
185
+ ttl: datetime = Field(..., description="Time to live")
186
+ created_at: datetime = Field(..., description="Creation timestamp")
187
+ approved_at: Optional[datetime] = Field(None, description="Approval timestamp")
188
+ approved_by: Optional[str] = Field(None, description="Approver email")
189
+ request_data: Dict[str, Any] = Field(default_factory=dict, description="Request data")
190
+ metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
191
+
192
+
193
+ class PolicyListResponse(BaseModel):
194
+ """Response for list policies endpoint"""
195
+ policies: List[Policy]
196
+ total: int
197
+ page: Optional[int] = 1 # Optional because enforcer may not return it
198
+ limit: int
199
+ has_more: bool
200
+
201
+
202
+ class RequestListResponse(BaseModel):
203
+ """Response for list requests endpoint"""
204
+ requests: List[ApprovalRequest]
205
+ total: int
206
+ page: int
207
+ limit: int
208
+ has_more: bool
209
+
210
+
211
+ # ============================================================================
212
+ # Specialized Clients (Separation of Concerns)
213
+ # ============================================================================
214
+
215
+ class PolicyOperations:
216
+ """Handles policy CRUD operations"""
217
+
218
+ def __init__(self, client: httpx.AsyncClient, base_url: str, headers: Dict[str, str]):
219
+ self._client = client
220
+ self._base_url = base_url
221
+ self._headers = headers
222
+
223
+ async def create(self, policy: PolicyCreate) -> Policy:
224
+ """
225
+ Create a new OPA policy.
226
+
227
+ Args:
228
+ policy: Policy creation data
229
+
230
+ Returns:
231
+ Created Policy object
232
+
233
+ Raises:
234
+ PolicyValidationError: If policy content is invalid
235
+ EnforcerConnectionError: If connection fails
236
+ """
237
+ try:
238
+ url = f"{self._base_url}/api/v1/policies"
239
+ response = await self._client.post(
240
+ url,
241
+ json=policy.model_dump(exclude_none=True),
242
+ headers=self._headers
243
+ )
244
+
245
+ if response.status_code == 201:
246
+ data = response.json()
247
+ logger.info(
248
+ "policy_created",
249
+ policy_id=data.get("id"),
250
+ policy_name=policy.name,
251
+ )
252
+ return Policy(**data)
253
+ elif response.status_code == 400:
254
+ error_data = response.json()
255
+ error_message = error_data.get("error", "Policy validation failed")
256
+ error_details = error_data.get("details", {})
257
+ # Handle both single error and array of errors
258
+ error_list = []
259
+ if "errors" in error_data and isinstance(error_data["errors"], list):
260
+ error_list = error_data["errors"]
261
+ elif error_details and "reason" in error_details:
262
+ error_list = [error_details["reason"]]
263
+
264
+ logger.error(
265
+ "policy_validation_failed",
266
+ error=error_message,
267
+ details=error_details,
268
+ errors=error_list
269
+ )
270
+ raise PolicyValidationError(
271
+ message=error_message,
272
+ details=error_details,
273
+ errors=error_list
274
+ )
275
+ elif response.status_code == 401:
276
+ raise EnforcerAuthenticationError("Authentication failed")
277
+ elif response.status_code == 409:
278
+ error_data = response.json()
279
+ raise PolicyValidationError(
280
+ message=error_data.get("error", "Policy already exists"),
281
+ details=error_data.get("details", {}),
282
+ errors=error_data.get("errors", [])
283
+ )
284
+ else:
285
+ raise PolicyEnforcerError(
286
+ f"Failed to create policy: HTTP {response.status_code}"
287
+ )
288
+
289
+ except httpx.RequestError as e:
290
+ logger.error("policy_creation_request_failed", error=str(e))
291
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
292
+
293
+ async def get(self, policy_id: str) -> Policy:
294
+ """
295
+ Get a specific policy by ID.
296
+
297
+ Args:
298
+ policy_id: Policy UUID
299
+
300
+ Returns:
301
+ Policy object
302
+
303
+ Raises:
304
+ PolicyNotFoundError: If policy doesn't exist
305
+ EnforcerConnectionError: If connection fails
306
+ """
307
+ try:
308
+ url = f"{self._base_url}/api/v1/policies/{policy_id}"
309
+ response = await self._client.get(url, headers=self._headers)
310
+
311
+ if response.status_code == 200:
312
+ return Policy(**response.json())
313
+ elif response.status_code == 404:
314
+ raise PolicyNotFoundError(f"Policy {policy_id} not found")
315
+ elif response.status_code == 401:
316
+ raise EnforcerAuthenticationError("Authentication failed")
317
+ else:
318
+ raise PolicyEnforcerError(
319
+ f"Failed to get policy: HTTP {response.status_code}"
320
+ )
321
+
322
+ except httpx.RequestError as e:
323
+ logger.error("policy_get_request_failed", error=str(e))
324
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
325
+
326
+ async def list(
327
+ self,
328
+ page: int = 1,
329
+ limit: int = 20,
330
+ enabled: Optional[bool] = None,
331
+ search: Optional[str] = None,
332
+ ) -> PolicyListResponse:
333
+ """
334
+ List policies with pagination and filtering.
335
+
336
+ Args:
337
+ page: Page number (default: 1)
338
+ limit: Items per page (default: 20, max: 100)
339
+ enabled: Filter by enabled status
340
+ search: Search term for policy name or description
341
+
342
+ Returns:
343
+ PolicyListResponse with policies and pagination info
344
+ """
345
+ try:
346
+ url = f"{self._base_url}/api/v1/policies"
347
+ params = {"page": page, "limit": min(limit, 100)}
348
+
349
+ if enabled is not None:
350
+ params["enabled"] = enabled
351
+ if search:
352
+ params["search"] = search
353
+
354
+ response = await self._client.get(
355
+ url,
356
+ params=params,
357
+ headers=self._headers
358
+ )
359
+
360
+ if response.status_code == 200:
361
+ data = response.json()
362
+ logger.info(
363
+ "policies_listed",
364
+ count=len(data.get("policies", [])),
365
+ total=data.get("total", 0),
366
+ )
367
+ return PolicyListResponse(**data)
368
+ elif response.status_code == 401:
369
+ raise EnforcerAuthenticationError("Authentication failed")
370
+ else:
371
+ raise PolicyEnforcerError(
372
+ f"Failed to list policies: HTTP {response.status_code}"
373
+ )
374
+
375
+ except httpx.RequestError as e:
376
+ logger.error("policies_list_request_failed", error=str(e))
377
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
378
+
379
+ async def update(self, policy_id: str, update: PolicyUpdate) -> Policy:
380
+ """
381
+ Update an existing policy.
382
+
383
+ Args:
384
+ policy_id: Policy UUID
385
+ update: Policy update data
386
+
387
+ Returns:
388
+ Updated Policy object
389
+
390
+ Raises:
391
+ PolicyNotFoundError: If policy doesn't exist
392
+ PolicyValidationError: If update is invalid
393
+ """
394
+ try:
395
+ url = f"{self._base_url}/api/v1/policies/{policy_id}"
396
+ response = await self._client.put(
397
+ url,
398
+ json=update.model_dump(exclude_none=True),
399
+ headers=self._headers
400
+ )
401
+
402
+ if response.status_code == 200:
403
+ data = response.json()
404
+ logger.info("policy_updated", policy_id=policy_id)
405
+ return Policy(**data)
406
+ elif response.status_code == 404:
407
+ raise PolicyNotFoundError(f"Policy {policy_id} not found")
408
+ elif response.status_code == 400:
409
+ error_data = response.json()
410
+ raise PolicyValidationError(
411
+ error_data.get("error", "Policy validation failed")
412
+ )
413
+ elif response.status_code == 401:
414
+ raise EnforcerAuthenticationError("Authentication failed")
415
+ else:
416
+ raise PolicyEnforcerError(
417
+ f"Failed to update policy: HTTP {response.status_code}"
418
+ )
419
+
420
+ except httpx.RequestError as e:
421
+ logger.error("policy_update_request_failed", error=str(e))
422
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
423
+
424
+ async def delete(self, policy_id: str) -> None:
425
+ """
426
+ Delete a policy.
427
+
428
+ Args:
429
+ policy_id: Policy UUID
430
+
431
+ Raises:
432
+ PolicyNotFoundError: If policy doesn't exist
433
+ """
434
+ try:
435
+ url = f"{self._base_url}/api/v1/policies/{policy_id}"
436
+ response = await self._client.delete(url, headers=self._headers)
437
+
438
+ if response.status_code == 204:
439
+ logger.info("policy_deleted", policy_id=policy_id)
440
+ return
441
+ elif response.status_code == 404:
442
+ raise PolicyNotFoundError(f"Policy {policy_id} not found")
443
+ elif response.status_code == 401:
444
+ raise EnforcerAuthenticationError("Authentication failed")
445
+ else:
446
+ raise PolicyEnforcerError(
447
+ f"Failed to delete policy: HTTP {response.status_code}"
448
+ )
449
+
450
+ except httpx.RequestError as e:
451
+ logger.error("policy_delete_request_failed", error=str(e))
452
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
453
+
454
+ async def validate(self, policy_id: str) -> ValidationResult:
455
+ """
456
+ Validate a policy's syntax and structure.
457
+
458
+ Args:
459
+ policy_id: Policy UUID
460
+
461
+ Returns:
462
+ ValidationResult with validation status and errors
463
+
464
+ Raises:
465
+ PolicyNotFoundError: If policy doesn't exist
466
+ """
467
+ try:
468
+ url = f"{self._base_url}/api/v1/policies/{policy_id}/validate"
469
+ response = await self._client.post(url, headers=self._headers)
470
+
471
+ if response.status_code in (200, 400):
472
+ data = response.json()
473
+ result = ValidationResult(**data)
474
+ logger.info(
475
+ "policy_validated",
476
+ policy_id=policy_id,
477
+ valid=result.valid,
478
+ error_count=len(result.errors),
479
+ )
480
+ return result
481
+ elif response.status_code == 404:
482
+ raise PolicyNotFoundError(f"Policy {policy_id} not found")
483
+ elif response.status_code == 401:
484
+ raise EnforcerAuthenticationError("Authentication failed")
485
+ else:
486
+ raise PolicyEnforcerError(
487
+ f"Failed to validate policy: HTTP {response.status_code}"
488
+ )
489
+
490
+ except httpx.RequestError as e:
491
+ logger.error("policy_validate_request_failed", error=str(e))
492
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
493
+
494
+
495
+ class EvaluationOperations:
496
+ """Handles policy evaluation operations"""
497
+
498
+ def __init__(self, client: httpx.AsyncClient, base_url: str, headers: Dict[str, str]):
499
+ self._client = client
500
+ self._base_url = base_url
501
+ self._headers = headers
502
+
503
+ async def evaluate(
504
+ self,
505
+ input_data: Dict[str, Any],
506
+ policy_id: Optional[str] = None,
507
+ policy_name: Optional[str] = None,
508
+ ) -> EvaluationResult:
509
+ """
510
+ Evaluate a policy against input data.
511
+
512
+ Args:
513
+ input_data: Input data to evaluate
514
+ policy_id: Policy UUID (use this or policy_name)
515
+ policy_name: Policy name (use this or policy_id)
516
+
517
+ Returns:
518
+ EvaluationResult with decision and violations
519
+
520
+ Raises:
521
+ PolicyEvaluationError: If evaluation fails
522
+ PolicyNotFoundError: If policy doesn't exist
523
+ """
524
+ try:
525
+ url = f"{self._base_url}/api/v1/evaluate"
526
+ request = EvaluationRequest(
527
+ input=input_data,
528
+ policy_id=policy_id,
529
+ policy_name=policy_name
530
+ )
531
+
532
+ response = await self._client.post(
533
+ url,
534
+ json=request.model_dump(exclude_none=True),
535
+ headers=self._headers
536
+ )
537
+
538
+ if response.status_code == 200:
539
+ result = EvaluationResult(**response.json())
540
+ logger.info(
541
+ "policy_evaluated",
542
+ policy_id=policy_id,
543
+ policy_name=policy_name,
544
+ allow=result.allow,
545
+ decision=result.decision,
546
+ violations=len(result.violations),
547
+ )
548
+ return result
549
+ elif response.status_code == 404:
550
+ raise PolicyNotFoundError(
551
+ f"Policy {policy_id or policy_name} not found"
552
+ )
553
+ elif response.status_code == 400:
554
+ error_data = response.json()
555
+ raise PolicyEvaluationError(
556
+ error_data.get("error", "Evaluation failed")
557
+ )
558
+ elif response.status_code == 401:
559
+ raise EnforcerAuthenticationError("Authentication failed")
560
+ else:
561
+ raise PolicyEvaluationError(
562
+ f"Evaluation failed: HTTP {response.status_code}"
563
+ )
564
+
565
+ except httpx.RequestError as e:
566
+ logger.error("evaluation_request_failed", error=str(e))
567
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
568
+
569
+ async def enforce(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
570
+ """
571
+ Call the enforcer's /enforce endpoint to evaluate ALL loaded policies.
572
+
573
+ This is the main enforcement endpoint that evaluates the input against
574
+ all policies loaded in the enforcer service.
575
+
576
+ Args:
577
+ input_data: Input data to evaluate against policies
578
+
579
+ Returns:
580
+ Dict with enforcement result:
581
+ {
582
+ "id": "enforcement-uuid",
583
+ "allow": true/false,
584
+ "policies": ["policy.names.that.passed.or.blocked"]
585
+ }
586
+
587
+ Raises:
588
+ PolicyEvaluationError: If enforcement fails
589
+ EnforcerConnectionError: If connection fails
590
+ """
591
+ try:
592
+ url = f"{self._base_url}/enforce"
593
+ response = await self._client.post(
594
+ url,
595
+ json=input_data,
596
+ headers=self._headers
597
+ )
598
+
599
+ if response.status_code == 200:
600
+ result = response.json()
601
+ logger.info(
602
+ "policy_enforced",
603
+ allow=result.get("allow"),
604
+ policies=result.get("policies", []),
605
+ enforcement_id=result.get("id"),
606
+ )
607
+ return result
608
+ elif response.status_code == 401:
609
+ raise EnforcerAuthenticationError("Authentication failed")
610
+ else:
611
+ raise PolicyEvaluationError(
612
+ f"Enforcement failed: HTTP {response.status_code}"
613
+ )
614
+
615
+ except httpx.RequestError as e:
616
+ logger.error("enforcement_request_failed", error=str(e))
617
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
618
+
619
+
620
+ class RequestOperations:
621
+ """Handles approval request operations"""
622
+
623
+ def __init__(self, client: httpx.AsyncClient, base_url: str, headers: Dict[str, str]):
624
+ self._client = client
625
+ self._base_url = base_url
626
+ self._headers = headers
627
+
628
+ async def approve(
629
+ self,
630
+ request_id: str,
631
+ metadata: Optional[Dict[str, Any]] = None,
632
+ ) -> ApprovalRequest:
633
+ """
634
+ Approve a pending request.
635
+
636
+ Args:
637
+ request_id: Request ID to approve
638
+ metadata: Optional metadata to attach
639
+
640
+ Returns:
641
+ Updated ApprovalRequest
642
+
643
+ Raises:
644
+ RequestNotFoundError: If request doesn't exist
645
+ """
646
+ try:
647
+ url = f"{self._base_url}/api/v1/requests/{request_id}/approve"
648
+ payload = {"metadata": metadata or {}}
649
+
650
+ response = await self._client.post(
651
+ url,
652
+ json=payload,
653
+ headers=self._headers
654
+ )
655
+
656
+ if response.status_code == 200:
657
+ logger.info("request_approved", request_id=request_id)
658
+ return ApprovalRequest(**response.json())
659
+ elif response.status_code == 404:
660
+ raise RequestNotFoundError(f"Request {request_id} not found")
661
+ elif response.status_code == 409:
662
+ error_data = response.json()
663
+ raise PolicyEnforcerError(
664
+ error_data.get("error", "Request already processed")
665
+ )
666
+ elif response.status_code == 401:
667
+ raise EnforcerAuthenticationError("Authentication failed")
668
+ else:
669
+ raise PolicyEnforcerError(
670
+ f"Failed to approve request: HTTP {response.status_code}"
671
+ )
672
+
673
+ except httpx.RequestError as e:
674
+ logger.error("approve_request_failed", error=str(e))
675
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
676
+
677
+ async def bulk_approve(
678
+ self,
679
+ request_ids: List[str],
680
+ metadata: Optional[Dict[str, Any]] = None,
681
+ ) -> Dict[str, Any]:
682
+ """
683
+ Approve multiple requests at once.
684
+
685
+ Args:
686
+ request_ids: List of request IDs (max 100)
687
+ metadata: Optional metadata to attach
688
+
689
+ Returns:
690
+ Dict with approval results
691
+ """
692
+ try:
693
+ url = f"{self._base_url}/api/v1/requests/bulk-approve"
694
+ payload = {
695
+ "request_ids": request_ids[:100],
696
+ "metadata": metadata or {},
697
+ }
698
+
699
+ response = await self._client.post(
700
+ url,
701
+ json=payload,
702
+ headers=self._headers
703
+ )
704
+
705
+ if response.status_code == 200:
706
+ data = response.json()
707
+ logger.info(
708
+ "requests_bulk_approved",
709
+ requested=data.get("requested_count", 0),
710
+ approved=data.get("approved_count", 0),
711
+ )
712
+ return data
713
+ elif response.status_code == 401:
714
+ raise EnforcerAuthenticationError("Authentication failed")
715
+ else:
716
+ raise PolicyEnforcerError(
717
+ f"Bulk approval failed: HTTP {response.status_code}"
718
+ )
719
+
720
+ except httpx.RequestError as e:
721
+ logger.error("bulk_approve_request_failed", error=str(e))
722
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
723
+
724
+ async def get(self, request_id: str) -> ApprovalRequest:
725
+ """
726
+ Get details of a specific request.
727
+
728
+ Args:
729
+ request_id: Request ID
730
+
731
+ Returns:
732
+ ApprovalRequest object
733
+
734
+ Raises:
735
+ RequestNotFoundError: If request doesn't exist
736
+ """
737
+ try:
738
+ url = f"{self._base_url}/api/v1/requests/{request_id}/describe"
739
+ response = await self._client.get(url, headers=self._headers)
740
+
741
+ if response.status_code == 200:
742
+ return ApprovalRequest(**response.json())
743
+ elif response.status_code == 404:
744
+ raise RequestNotFoundError(f"Request {request_id} not found")
745
+ elif response.status_code == 401:
746
+ raise EnforcerAuthenticationError("Authentication failed")
747
+ else:
748
+ raise PolicyEnforcerError(
749
+ f"Failed to get request: HTTP {response.status_code}"
750
+ )
751
+
752
+ except httpx.RequestError as e:
753
+ logger.error("get_request_failed", error=str(e))
754
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
755
+
756
+ async def list(
757
+ self,
758
+ page: int = 1,
759
+ limit: int = 20,
760
+ status: Optional[RequestStatus] = None,
761
+ runner: Optional[str] = None,
762
+ ) -> RequestListResponse:
763
+ """
764
+ List approval requests with pagination.
765
+
766
+ Args:
767
+ page: Page number
768
+ limit: Items per page (max 100)
769
+ status: Filter by status
770
+ runner: Filter by runner name
771
+
772
+ Returns:
773
+ RequestListResponse with requests and pagination
774
+ """
775
+ try:
776
+ url = f"{self._base_url}/api/v1/requests"
777
+ params = {"page": page, "limit": min(limit, 100)}
778
+
779
+ if status:
780
+ params["status"] = status.value
781
+ if runner:
782
+ params["runner"] = runner
783
+
784
+ response = await self._client.get(
785
+ url,
786
+ params=params,
787
+ headers=self._headers
788
+ )
789
+
790
+ if response.status_code == 200:
791
+ data = response.json()
792
+ logger.info(
793
+ "requests_listed",
794
+ count=len(data.get("requests", [])),
795
+ total=data.get("total", 0),
796
+ )
797
+ return RequestListResponse(**data)
798
+ elif response.status_code == 401:
799
+ raise EnforcerAuthenticationError("Authentication failed")
800
+ else:
801
+ raise PolicyEnforcerError(
802
+ f"Failed to list requests: HTTP {response.status_code}"
803
+ )
804
+
805
+ except httpx.RequestError as e:
806
+ logger.error("list_requests_failed", error=str(e))
807
+ raise EnforcerConnectionError(f"Connection failed: {str(e)}") from e
808
+
809
+
810
+ # ============================================================================
811
+ # Main Client with Context Manager Support
812
+ # ============================================================================
813
+
814
+ class PolicyEnforcerClient:
815
+ """
816
+ Main client for OPA Watchdog Enforcer Service.
817
+
818
+ This client provides a clean, async interface with proper resource management.
819
+ Use it as a context manager to ensure proper cleanup:
820
+
821
+ async with PolicyEnforcerClient(base_url="...", api_key="...") as client:
822
+ policy = await client.policies.create(...)
823
+ result = await client.evaluation.evaluate(...)
824
+ request = await client.requests.approve(...)
825
+
826
+ Attributes:
827
+ policies: PolicyOperations for CRUD operations
828
+ evaluation: EvaluationOperations for policy evaluation
829
+ requests: RequestOperations for approval workflows
830
+ """
831
+
832
+ def __init__(
833
+ self,
834
+ base_url: str,
835
+ api_key: str,
836
+ timeout: float = 30.0,
837
+ max_retries: int = 3,
838
+ auth_type: str = "UserKey",
839
+ ):
840
+ """
841
+ Initialize Policy Enforcer client.
842
+
843
+ Args:
844
+ base_url: Enforcer service URL
845
+ api_key: Kubiya API key (JWT token)
846
+ timeout: Request timeout in seconds
847
+ max_retries: Maximum number of retries for failed requests
848
+ auth_type: Authentication type - "UserKey" or "Bearer" (default: "UserKey")
849
+ """
850
+ self._base_url = base_url.rstrip("/")
851
+ self._api_key = api_key
852
+ self._timeout = timeout
853
+ self._max_retries = max_retries
854
+ self._auth_type = auth_type
855
+ self._headers = {"Authorization": f"{auth_type} {api_key}"}
856
+
857
+ # Create async HTTP client with retries
858
+ transport = httpx.AsyncHTTPTransport(retries=max_retries)
859
+ self._client = httpx.AsyncClient(
860
+ timeout=httpx.Timeout(timeout, connect=5.0),
861
+ limits=httpx.Limits(max_connections=20, max_keepalive_connections=10),
862
+ transport=transport,
863
+ )
864
+
865
+ # Initialize specialized operation handlers
866
+ self.policies = PolicyOperations(self._client, self._base_url, self._headers)
867
+ self.evaluation = EvaluationOperations(self._client, self._base_url, self._headers)
868
+ self.requests = RequestOperations(self._client, self._base_url, self._headers)
869
+
870
+ async def __aenter__(self):
871
+ """Context manager entry"""
872
+ return self
873
+
874
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
875
+ """Context manager exit with cleanup"""
876
+ await self.close()
877
+
878
+ async def close(self):
879
+ """Close the HTTP client and cleanup resources"""
880
+ try:
881
+ await self._client.aclose()
882
+ except Exception as e:
883
+ logger.warning("client_close_error", error=str(e), error_type=type(e).__name__)
884
+
885
+ async def health_check(self) -> bool:
886
+ """
887
+ Check if the enforcer service is healthy.
888
+
889
+ Returns:
890
+ True if healthy, False otherwise
891
+ """
892
+ try:
893
+ # Use the policies list endpoint with limit=1 to check connectivity
894
+ url = f"{self._base_url}/api/v1/policies"
895
+ response = await self._client.get(
896
+ url,
897
+ params={"limit": 1},
898
+ headers=self._headers,
899
+ timeout=5.0
900
+ )
901
+ return response.status_code == 200
902
+ except Exception as e:
903
+ logger.warning("health_check_failed", error=str(e), error_type=type(e).__name__)
904
+ return False
905
+
906
+ async def get_status(self) -> Dict[str, Any]:
907
+ """
908
+ Get detailed service status.
909
+
910
+ Returns:
911
+ Status dict with service information
912
+ """
913
+ try:
914
+ url = f"{self._base_url}/status"
915
+ response = await self._client.get(url, timeout=5.0)
916
+ if response.status_code == 200:
917
+ return response.json()
918
+ return {}
919
+ except Exception as e:
920
+ logger.warning("status_check_failed", error=str(e))
921
+ return {}
922
+
923
+
924
+ # ============================================================================
925
+ # Factory Function (Dependency Injection)
926
+ # ============================================================================
927
+
928
+ @asynccontextmanager
929
+ async def create_policy_enforcer_client(
930
+ enforcer_url: Optional[str] = None,
931
+ api_key: Optional[str] = None,
932
+ auth_type: str = "UserKey",
933
+ ) -> Optional[PolicyEnforcerClient]:
934
+ """
935
+ Factory function to create a PolicyEnforcerClient with context manager support.
936
+
937
+ Reads configuration from environment variables if not provided:
938
+ - ENFORCER_SERVICE_URL: Enforcer service URL (default: https://enforcer-psi.vercel.app)
939
+ - api_key: Authorization token (typically passed from the incoming request)
940
+
941
+ Args:
942
+ enforcer_url: Optional enforcer URL override
943
+ api_key: Authorization token (Bearer token from the request)
944
+ auth_type: Authentication type - "UserKey" or "Bearer" (default: "UserKey")
945
+
946
+ Yields:
947
+ PolicyEnforcerClient instance if configured, None if disabled
948
+
949
+ Usage:
950
+ async with create_policy_enforcer_client(api_key=request_token, auth_type="UserKey") as client:
951
+ if client:
952
+ policy = await client.policies.create(...)
953
+ """
954
+ # Check if enforcer is enabled
955
+ enforcer_url = enforcer_url or os.environ.get("ENFORCER_SERVICE_URL")
956
+
957
+ # If no URL is set, yield None (enforcer is disabled)
958
+ if not enforcer_url:
959
+ logger.info("policy_enforcer_disabled", reason="no_url")
960
+ yield None
961
+ return
962
+
963
+ # Strip whitespace and newlines
964
+ enforcer_url = enforcer_url.strip()
965
+
966
+ # Default to production URL
967
+ if enforcer_url == "":
968
+ enforcer_url = "https://enforcer-psi.vercel.app"
969
+
970
+ # API key should be passed from the request, not from environment
971
+ # Fall back to KUBIYA_API_KEY for backward compatibility
972
+ if not api_key:
973
+ api_key = os.environ.get("KUBIYA_API_KEY")
974
+
975
+ if not api_key:
976
+ logger.warning("policy_enforcer_disabled_no_api_key", reason="missing_token")
977
+ yield None
978
+ return
979
+
980
+ # Create and yield client
981
+ client = PolicyEnforcerClient(base_url=enforcer_url, api_key=api_key, auth_type=auth_type)
982
+
983
+ try:
984
+ logger.info("policy_enforcer_client_created", enforcer_url=enforcer_url, auth_type=auth_type)
985
+ yield client
986
+ finally:
987
+ try:
988
+ await client.close()
989
+ except Exception as e:
990
+ logger.warning("policy_enforcer_client_cleanup_error", error=str(e), error_type=type(e).__name__)
991
+
992
+
993
+ # Convenience function for dependency injection in FastAPI
994
+ def get_policy_enforcer_client_dependency() -> Optional[PolicyEnforcerClient]:
995
+ """
996
+ Dependency function for FastAPI to inject PolicyEnforcerClient.
997
+
998
+ Usage in FastAPI:
999
+ @router.get("/policies")
1000
+ async def list_policies(
1001
+ client: PolicyEnforcerClient = Depends(get_policy_enforcer_client_dependency)
1002
+ ):
1003
+ if client:
1004
+ policies = await client.policies.list()
1005
+ ...
1006
+ """
1007
+ enforcer_url = os.environ.get("ENFORCER_SERVICE_URL")
1008
+ api_key = os.environ.get("KUBIYA_API_KEY")
1009
+
1010
+ if not enforcer_url or not api_key:
1011
+ return None
1012
+
1013
+ if enforcer_url == "":
1014
+ enforcer_url = "https://enforcer-psi.vercel.app"
1015
+
1016
+ return PolicyEnforcerClient(base_url=enforcer_url, api_key=api_key)