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,979 @@
1
+ """
2
+ Base runtime abstraction with proper ABC, registry, and Control Plane integration.
3
+
4
+ This module provides:
5
+ - Abstract base class for all runtimes
6
+ - Runtime registry for discovery and instantiation
7
+ - Lifecycle hooks for extensibility
8
+ - Control Plane integration helpers
9
+ - Configuration validation framework
10
+ """
11
+
12
+ from typing import AsyncIterator, Dict, Any, Optional, List, Callable, Type
13
+ from dataclasses import dataclass, field
14
+ from enum import Enum
15
+ from abc import ABC, abstractmethod
16
+ import structlog
17
+
18
+ logger = structlog.get_logger(__name__)
19
+
20
+
21
+ class RuntimeType(str, Enum):
22
+ """Enumeration of supported runtime types."""
23
+
24
+ DEFAULT = "default" # Agno-based runtime
25
+ CLAUDE_CODE = "claude_code" # Claude Code SDK runtime
26
+ AGENT_RUNTIME = "agent_runtime" # Rust-based high-performance runtime
27
+ # Easy to add more: LANGCHAIN = "langchain", CREWAI = "crewai", etc.
28
+
29
+
30
+ @dataclass
31
+ class RuntimeExecutionResult:
32
+ """
33
+ Standardized result structure from any runtime.
34
+
35
+ This ensures all runtimes return consistent data structures that can
36
+ be consumed by the workflow and activity layers.
37
+
38
+ Analytics Integration:
39
+ The `usage` field provides standardized token usage metrics that are
40
+ automatically extracted and submitted to the analytics system.
41
+ """
42
+
43
+ response: str
44
+ """The main response text from the agent."""
45
+
46
+ usage: Dict[str, Any]
47
+ """
48
+ Token usage metrics with standardized fields:
49
+ - input_tokens (int): Number of input/prompt tokens
50
+ - output_tokens (int): Number of output/completion tokens
51
+ - total_tokens (int): Total tokens used
52
+ - cache_read_tokens (int, optional): Cached tokens read (Anthropic)
53
+ - cache_creation_tokens (int, optional): Tokens used for cache creation (Anthropic)
54
+ - prompt_tokens_details (dict, optional): Detailed breakdown from provider
55
+
56
+ Runtimes should populate this from their native usage tracking.
57
+ """
58
+
59
+ success: bool
60
+ """Whether the execution succeeded."""
61
+
62
+ finish_reason: Optional[str] = None
63
+ """Reason the execution finished (e.g., 'stop', 'length', 'tool_use')."""
64
+
65
+ run_id: Optional[str] = None
66
+ """Unique identifier for this execution run."""
67
+
68
+ model: Optional[str] = None
69
+ """Model identifier used for this execution."""
70
+
71
+ tool_execution_messages: Optional[List[Dict]] = None
72
+ """
73
+ Tool execution messages for UI display and analytics.
74
+ Format: [{"tool": "Bash", "input": {...}, "output": {...}, "success": bool, "duration_ms": int}, ...]
75
+
76
+ Analytics Integration:
77
+ These are automatically tracked in the execution_tool_calls table.
78
+ """
79
+
80
+ tool_messages: Optional[List[Dict]] = None
81
+ """
82
+ Detailed tool messages with execution metadata.
83
+ Format: [{"role": "tool", "content": "...", "tool_use_id": "..."}, ...]
84
+ """
85
+
86
+ error: Optional[str] = None
87
+ """Error message if execution failed."""
88
+
89
+ metadata: Dict[str, Any] = field(default_factory=dict)
90
+ """
91
+ Additional runtime-specific metadata.
92
+
93
+ Can include:
94
+ - turn_duration_ms (int): Duration of this turn in milliseconds
95
+ - model_provider (str): Provider name (anthropic, openai, google, etc.)
96
+ - tasks (list): Task tracking data for analytics
97
+ - any runtime-specific metrics
98
+ """
99
+
100
+
101
+ @dataclass
102
+ class RuntimeExecutionContext:
103
+ """
104
+ Context passed to runtime for execution.
105
+
106
+ This contains all the information needed for an agent to execute,
107
+ regardless of which runtime implementation is used.
108
+ """
109
+
110
+ execution_id: str
111
+ """Unique identifier for this execution."""
112
+
113
+ agent_id: str
114
+ """Unique identifier for the agent being executed."""
115
+
116
+ organization_id: str
117
+ """Organization context for this execution."""
118
+
119
+ prompt: str
120
+ """User's input prompt/message."""
121
+
122
+ system_prompt: Optional[str] = None
123
+ """System-level instructions for the agent."""
124
+
125
+ conversation_history: List[Dict[str, Any]] = field(default_factory=list)
126
+ """
127
+ Previous conversation messages.
128
+ Format: [{"role": "user|assistant|system", "content": "..."}, ...]
129
+ """
130
+
131
+ model_id: Optional[str] = None
132
+ """LiteLLM model identifier (e.g., 'gpt-4', 'claude-3-opus')."""
133
+
134
+ model_config: Optional[Dict[str, Any]] = None
135
+ """Model-specific configuration (temperature, top_p, etc.)."""
136
+
137
+ agent_config: Optional[Dict[str, Any]] = None
138
+ """Agent-specific configuration."""
139
+
140
+ skills: List[Any] = field(default_factory=list)
141
+ """Resolved skills available to the agent."""
142
+
143
+ skill_configs: List[Dict[str, Any]] = field(default_factory=list)
144
+ """Original skill configuration dictionaries (before instantiation)."""
145
+
146
+ mcp_servers: Optional[Dict[str, Any]] = None
147
+ """MCP server configurations."""
148
+
149
+ user_metadata: Optional[Dict[str, Any]] = None
150
+ """User-provided metadata for this execution."""
151
+
152
+ runtime_config: Optional[Dict[str, Any]] = None
153
+ """Runtime-specific configuration options."""
154
+
155
+ runtime_type: Optional[RuntimeType] = None
156
+ """Runtime type for this execution."""
157
+
158
+ # Enforcement context fields
159
+ user_email: Optional[str] = None
160
+ """User email for enforcement context."""
161
+
162
+ user_id: Optional[str] = None
163
+ """User ID for enforcement context."""
164
+
165
+ user_roles: List[str] = field(default_factory=list)
166
+ """User roles for enforcement context."""
167
+
168
+ team_id: Optional[str] = None
169
+ """Team ID for enforcement context."""
170
+
171
+ team_name: Optional[str] = None
172
+ """Team name for enforcement context."""
173
+
174
+ environment: str = "production"
175
+ """Environment for enforcement context (dev/staging/production)."""
176
+
177
+ workspace_directory: Optional[str] = None
178
+ """
179
+ Execution workspace directory (e.g., .kubiya/workspaces/<execution-id>).
180
+ Used by runtimes and skills for file operations.
181
+ """
182
+
183
+ # ==================== Session Persistence Support ====================
184
+
185
+ session_id: Optional[str] = None
186
+ """Session identifier for multi-turn persistence."""
187
+
188
+ session_messages: List[Dict[str, Any]] = field(default_factory=list)
189
+ """
190
+ Full session messages with metadata (not just conversation_history).
191
+
192
+ Format: [{
193
+ 'role': str, # "user", "assistant", "system", "tool"
194
+ 'content': str,
195
+ 'timestamp': str, # ISO 8601 timestamp
196
+ 'message_id': str, # Deterministic message identifier
197
+ 'user_id': str,
198
+ 'user_name': str,
199
+ 'user_email': str,
200
+ 'user_avatar': str,
201
+ 'metadata': dict # Additional metadata (tool traces, etc.)
202
+ }]
203
+
204
+ Note: This differs from conversation_history which only has {role, content}.
205
+ session_messages includes full attribution and tracing information.
206
+ """
207
+
208
+ # ==================== Native Sub-Agent Execution Support ====================
209
+
210
+ agents: Optional[Dict[str, Dict[str, Any]]] = None
211
+ """
212
+ Sub-agent definitions for native execution.
213
+
214
+ Format: {
215
+ 'agent_id_1': {
216
+ 'description': str, # Agent role/purpose
217
+ 'prompt': str, # System prompt for the sub-agent
218
+ 'tools': List[str], # Available tool names
219
+ 'model': str, # "sonnet", "opus", "haiku", or "inherit"
220
+ 'config': dict # Optional runtime configuration
221
+ }
222
+ }
223
+
224
+ When enable_native_subagents=True, runtimes can use these definitions
225
+ to orchestrate sub-agents internally (similar to Claude Code's approach).
226
+ """
227
+
228
+ # ==================== Feature Flags ====================
229
+
230
+ enable_session_persistence: bool = False
231
+ """Whether runtime should persist sessions (opt-in)."""
232
+
233
+ enable_native_subagents: bool = False
234
+ """Whether to use native sub-agent execution (opt-in)."""
235
+
236
+
237
+ @dataclass
238
+ class RuntimeCapabilities:
239
+ """Runtime capabilities metadata."""
240
+
241
+ streaming: bool = False
242
+ """Supports streaming execution."""
243
+
244
+ tools: bool = False
245
+ """Supports tool calling."""
246
+
247
+ mcp: bool = False
248
+ """Supports MCP servers."""
249
+
250
+ hooks: bool = False
251
+ """Supports lifecycle hooks."""
252
+
253
+ cancellation: bool = False
254
+ """Supports execution cancellation."""
255
+
256
+ conversation_history: bool = False
257
+ """Supports multi-turn conversations."""
258
+
259
+ custom_tools: bool = False
260
+ """Supports custom tool registration."""
261
+
262
+
263
+ class BaseRuntime(ABC):
264
+ """
265
+ Abstract base class for all agent runtimes.
266
+
267
+ This class provides:
268
+ - Standard interface for all runtimes
269
+ - Lifecycle hooks for extensibility
270
+ - Control Plane integration helpers
271
+ - Configuration validation
272
+ - Error handling patterns
273
+
274
+ To create a new runtime:
275
+ 1. Inherit from BaseRuntime
276
+ 2. Implement abstract methods
277
+ 3. Register via @RuntimeRegistry.register()
278
+ 4. Override lifecycle hooks as needed
279
+ """
280
+
281
+ def __init__(
282
+ self,
283
+ control_plane_client: Any,
284
+ cancellation_manager: Any,
285
+ **kwargs,
286
+ ):
287
+ """
288
+ Initialize the runtime.
289
+
290
+ Args:
291
+ control_plane_client: Client for Control Plane API
292
+ cancellation_manager: Manager for execution cancellation
293
+ **kwargs: Additional configuration options
294
+ """
295
+ self.control_plane = control_plane_client
296
+ self.cancellation_manager = cancellation_manager
297
+ self.logger = structlog.get_logger(self.__class__.__name__)
298
+ self.config = kwargs
299
+
300
+ # Track active executions for cleanup
301
+ self._active_executions: Dict[str, Any] = {}
302
+
303
+ # ==================== Abstract Methods (Must Implement) ====================
304
+
305
+ @abstractmethod
306
+ async def _execute_impl(
307
+ self, context: RuntimeExecutionContext
308
+ ) -> RuntimeExecutionResult:
309
+ """
310
+ Core execution logic (non-streaming).
311
+
312
+ Implement the actual execution logic here without worrying about
313
+ lifecycle hooks, Control Plane integration, or error handling.
314
+ The base class handles those concerns.
315
+
316
+ Args:
317
+ context: Execution context
318
+
319
+ Returns:
320
+ RuntimeExecutionResult
321
+ """
322
+ pass
323
+
324
+ @abstractmethod
325
+ async def _stream_execute_impl(
326
+ self,
327
+ context: RuntimeExecutionContext,
328
+ event_callback: Optional[Callable[[Dict], None]] = None,
329
+ ) -> AsyncIterator[RuntimeExecutionResult]:
330
+ """
331
+ Core streaming execution logic.
332
+
333
+ Implement the actual streaming logic here. The base class
334
+ handles lifecycle hooks and error handling.
335
+
336
+ Args:
337
+ context: Execution context
338
+ event_callback: Optional callback for events
339
+
340
+ Yields:
341
+ RuntimeExecutionResult chunks
342
+ """
343
+ pass
344
+
345
+ @abstractmethod
346
+ def get_runtime_type(self) -> RuntimeType:
347
+ """Return the runtime type identifier."""
348
+ pass
349
+
350
+ @abstractmethod
351
+ def get_capabilities(self) -> RuntimeCapabilities:
352
+ """Return runtime capabilities."""
353
+ pass
354
+
355
+ # ==================== Public Interface (Don't Override) ====================
356
+
357
+ async def execute(
358
+ self, context: RuntimeExecutionContext
359
+ ) -> RuntimeExecutionResult:
360
+ """
361
+ Execute agent with full lifecycle management.
362
+
363
+ This method orchestrates the entire execution lifecycle:
364
+ 1. Validate configuration
365
+ 2. Call before_execute hook
366
+ 3. Cache metadata in Control Plane
367
+ 4. Execute via _execute_impl
368
+ 5. Call after_execute hook
369
+ 6. Handle errors via on_error hook
370
+ 7. Cleanup
371
+
372
+ Args:
373
+ context: Execution context
374
+
375
+ Returns:
376
+ RuntimeExecutionResult
377
+ """
378
+ execution_id = context.execution_id
379
+
380
+ try:
381
+ # 1. Validate configuration
382
+ self._validate_config(context)
383
+
384
+ # 2. Before execute hook
385
+ await self.before_execute(context)
386
+
387
+ # 3. Cache metadata in Control Plane
388
+ self._cache_execution_metadata(context)
389
+
390
+ # 4. Register for cancellation
391
+ if self.get_capabilities().cancellation:
392
+ self.cancellation_manager.register(
393
+ execution_id=execution_id,
394
+ instance=self,
395
+ instance_type=self.__class__.__name__,
396
+ )
397
+
398
+ # 5. Execute implementation
399
+ self.logger.info(
400
+ "🚀 Runtime execution started",
401
+ execution_id=execution_id,
402
+ runtime=self.get_runtime_type().value,
403
+ model=context.model_id or "default",
404
+ stream=False
405
+ )
406
+
407
+ result = await self._execute_impl(context)
408
+
409
+ # 6. After execute hook
410
+ await self.after_execute(context, result)
411
+
412
+ self.logger.info(
413
+ "runtime_execute_complete",
414
+ execution_id=execution_id[:8],
415
+ success=result.success,
416
+ )
417
+
418
+ return result
419
+
420
+ except Exception as e:
421
+ # 7. Error hook
422
+ error_result = await self.on_error(context, e)
423
+ return error_result
424
+
425
+ finally:
426
+ # 8. Cleanup
427
+ if self.get_capabilities().cancellation:
428
+ self.cancellation_manager.unregister(execution_id)
429
+ self._active_executions.pop(execution_id, None)
430
+
431
+ async def stream_execute(
432
+ self,
433
+ context: RuntimeExecutionContext,
434
+ event_callback: Optional[Callable[[Dict], None]] = None,
435
+ ) -> AsyncIterator[RuntimeExecutionResult]:
436
+ """
437
+ Execute agent with streaming and full lifecycle management.
438
+
439
+ Args:
440
+ context: Execution context
441
+ event_callback: Optional callback for events
442
+
443
+ Yields:
444
+ RuntimeExecutionResult chunks
445
+ """
446
+ execution_id = context.execution_id
447
+
448
+ try:
449
+ # 1. Validate configuration
450
+ self._validate_config(context)
451
+
452
+ # 2. Before execute hook
453
+ await self.before_execute(context)
454
+
455
+ # 3. Cache metadata in Control Plane
456
+ self._cache_execution_metadata(context)
457
+
458
+ # 4. Register for cancellation
459
+ if self.get_capabilities().cancellation:
460
+ self.cancellation_manager.register(
461
+ execution_id=execution_id,
462
+ instance=self,
463
+ instance_type=self.__class__.__name__,
464
+ )
465
+
466
+ # 5. Stream implementation
467
+ self.logger.info(
468
+ "🚀 Runtime streaming execution started",
469
+ execution_id=execution_id,
470
+ runtime=self.get_runtime_type().value,
471
+ model=context.model_id or "default",
472
+ stream=True
473
+ )
474
+
475
+ final_result = None
476
+ async for chunk in self._stream_execute_impl(context, event_callback):
477
+ yield chunk
478
+ if chunk.finish_reason:
479
+ final_result = chunk
480
+
481
+ # 6. After execute hook
482
+ if final_result:
483
+ await self.after_execute(context, final_result)
484
+
485
+ self.logger.info(
486
+ "runtime_stream_complete",
487
+ execution_id=execution_id[:8],
488
+ )
489
+
490
+ except Exception as e:
491
+ # 7. Error hook
492
+ error_result = await self.on_error(context, e)
493
+ yield error_result
494
+
495
+ finally:
496
+ # 8. Cleanup
497
+ if self.get_capabilities().cancellation:
498
+ self.cancellation_manager.unregister(execution_id)
499
+ self._active_executions.pop(execution_id, None)
500
+
501
+ async def cancel(self, execution_id: str) -> bool:
502
+ """
503
+ Cancel an in-progress execution.
504
+
505
+ Override _cancel_impl() to provide runtime-specific cancellation logic.
506
+
507
+ Args:
508
+ execution_id: ID of execution to cancel
509
+
510
+ Returns:
511
+ True if cancellation succeeded
512
+ """
513
+ if not self.get_capabilities().cancellation:
514
+ self.logger.warning(
515
+ "runtime_cancel_not_supported",
516
+ runtime=self.get_runtime_type().value,
517
+ )
518
+ return False
519
+
520
+ try:
521
+ return await self._cancel_impl(execution_id)
522
+ except Exception as e:
523
+ self.logger.error(
524
+ "runtime_cancel_failed",
525
+ execution_id=execution_id[:8],
526
+ error=str(e),
527
+ )
528
+ return False
529
+
530
+ async def get_usage(self, execution_id: str) -> Dict[str, Any]:
531
+ """
532
+ Get usage metrics for an execution.
533
+
534
+ Override _get_usage_impl() to provide runtime-specific usage tracking.
535
+
536
+ Args:
537
+ execution_id: ID of execution
538
+
539
+ Returns:
540
+ Usage metrics dict
541
+ """
542
+ try:
543
+ return await self._get_usage_impl(execution_id)
544
+ except Exception as e:
545
+ self.logger.error(
546
+ "runtime_get_usage_failed",
547
+ execution_id=execution_id[:8],
548
+ error=str(e),
549
+ )
550
+ return {}
551
+
552
+ # ==================== Capabilities API ====================
553
+
554
+ def supports_streaming(self) -> bool:
555
+ """Whether this runtime supports streaming execution."""
556
+ return self.get_capabilities().streaming
557
+
558
+ def supports_tools(self) -> bool:
559
+ """Whether this runtime supports tool calling."""
560
+ return self.get_capabilities().tools
561
+
562
+ def supports_mcp(self) -> bool:
563
+ """Whether this runtime supports MCP servers."""
564
+ return self.get_capabilities().mcp
565
+
566
+ def supports_custom_tools(self) -> bool:
567
+ """Whether this runtime supports custom tool extensions."""
568
+ return self.get_capabilities().custom_tools
569
+
570
+ def get_runtime_info(self) -> Dict[str, Any]:
571
+ """
572
+ Get information about this runtime implementation.
573
+
574
+ Override to provide additional metadata.
575
+
576
+ Returns:
577
+ Dict with runtime metadata
578
+ """
579
+ caps = self.get_capabilities()
580
+ return {
581
+ "runtime_type": self.get_runtime_type().value,
582
+ "supports_streaming": caps.streaming,
583
+ "supports_tools": caps.tools,
584
+ "supports_mcp": caps.mcp,
585
+ "supports_hooks": caps.hooks,
586
+ "supports_cancellation": caps.cancellation,
587
+ "supports_conversation_history": caps.conversation_history,
588
+ "supports_custom_tools": caps.custom_tools,
589
+ }
590
+
591
+ # ==================== Lifecycle Hooks (Override as Needed) ====================
592
+
593
+ async def before_execute(self, context: RuntimeExecutionContext) -> None:
594
+ """
595
+ Hook called before execution starts.
596
+
597
+ Override to:
598
+ - Validate additional configuration
599
+ - Setup resources
600
+ - Initialize connections
601
+ - Log execution start
602
+
603
+ Args:
604
+ context: Execution context
605
+ """
606
+ pass
607
+
608
+ async def after_execute(
609
+ self, context: RuntimeExecutionContext, result: RuntimeExecutionResult
610
+ ) -> None:
611
+ """
612
+ Hook called after successful execution.
613
+
614
+ Override to:
615
+ - Cleanup resources
616
+ - Log metrics
617
+ - Update statistics
618
+ - Trigger webhooks
619
+
620
+ Args:
621
+ context: Execution context
622
+ result: Execution result
623
+ """
624
+ pass
625
+
626
+ async def on_error(
627
+ self, context: RuntimeExecutionContext, error: Exception
628
+ ) -> RuntimeExecutionResult:
629
+ """
630
+ Hook called when execution fails.
631
+
632
+ Override to:
633
+ - Custom error handling
634
+ - Error reporting
635
+ - Cleanup
636
+ - Fallback logic
637
+
638
+ Args:
639
+ context: Execution context
640
+ error: Exception that occurred
641
+
642
+ Returns:
643
+ RuntimeExecutionResult with error details
644
+ """
645
+ self.logger.error(
646
+ "runtime_execution_failed",
647
+ execution_id=context.execution_id[:8],
648
+ runtime=self.get_runtime_type().value,
649
+ error=str(error),
650
+ error_type=type(error).__name__,
651
+ )
652
+
653
+ return RuntimeExecutionResult(
654
+ response="",
655
+ usage={},
656
+ success=False,
657
+ error=f"{type(error).__name__}: {str(error)}",
658
+ finish_reason="error",
659
+ )
660
+
661
+ # ==================== Helper Methods (Override as Needed) ====================
662
+
663
+ async def _cancel_impl(self, execution_id: str) -> bool:
664
+ """
665
+ Runtime-specific cancellation implementation.
666
+
667
+ Override to provide custom cancellation logic.
668
+
669
+ Args:
670
+ execution_id: ID of execution to cancel
671
+
672
+ Returns:
673
+ True if successful
674
+ """
675
+ return False
676
+
677
+ async def _get_usage_impl(self, execution_id: str) -> Dict[str, Any]:
678
+ """
679
+ Runtime-specific usage tracking implementation.
680
+
681
+ Override to provide usage metrics.
682
+
683
+ Args:
684
+ execution_id: ID of execution
685
+
686
+ Returns:
687
+ Usage metrics dict
688
+ """
689
+ return {}
690
+
691
+ def _validate_config(self, context: RuntimeExecutionContext) -> None:
692
+ """
693
+ Validate runtime configuration.
694
+
695
+ Override to add custom validation logic.
696
+ Raise ValueError if configuration is invalid.
697
+
698
+ Args:
699
+ context: Execution context
700
+
701
+ Raises:
702
+ ValueError: If configuration is invalid
703
+ """
704
+ # Base validation
705
+ if not context.prompt:
706
+ raise ValueError("Prompt is required")
707
+ if not context.execution_id:
708
+ raise ValueError("Execution ID is required")
709
+
710
+ # Runtime-specific requirements validation
711
+ try:
712
+ from control_plane_api.worker.runtimes.validation import RuntimeRequirementsRegistry
713
+
714
+ is_valid, errors = RuntimeRequirementsRegistry.validate_for_runtime(
715
+ self.get_runtime_type(), context
716
+ )
717
+
718
+ if not is_valid:
719
+ error_msg = "Runtime validation failed:\n" + "\n".join(f" - {err}" for err in errors)
720
+ raise ValueError(error_msg)
721
+
722
+ except ImportError:
723
+ # Validation module not available - skip validation
724
+ self.logger.warning("Runtime validation module not available, skipping validation")
725
+ except Exception as e:
726
+ self.logger.error(
727
+ "Runtime validation error",
728
+ error=str(e),
729
+ runtime=self.get_runtime_type().value,
730
+ )
731
+ raise
732
+
733
+ def _cache_execution_metadata(self, context: RuntimeExecutionContext) -> None:
734
+ """
735
+ Cache execution metadata in Control Plane.
736
+
737
+ This enables:
738
+ - Execution tracking
739
+ - Real-time monitoring
740
+ - Analytics
741
+
742
+ Args:
743
+ context: Execution context
744
+ """
745
+ try:
746
+ self.control_plane.cache_metadata(
747
+ context.execution_id,
748
+ "AGENT",
749
+ )
750
+ except Exception as e:
751
+ self.logger.warning(
752
+ "failed_to_cache_metadata",
753
+ execution_id=context.execution_id[:8],
754
+ error=str(e),
755
+ )
756
+
757
+ # ==================== Custom Tool Extension API ====================
758
+
759
+ def get_custom_tool_requirements(self) -> Dict[str, Any]:
760
+ """
761
+ Get requirements/documentation for creating custom tools for this runtime.
762
+
763
+ Override this method to document how developers should create custom tools
764
+ for your runtime.
765
+
766
+ Returns:
767
+ Dictionary with:
768
+ - format: Tool format (e.g., "python_class", "mcp_server")
769
+ - description: Human-readable description
770
+ - example_code: Example implementation
771
+ - documentation_url: Link to detailed docs
772
+ - required_methods: List of required methods/attributes
773
+ - schema: JSON schema for validation
774
+
775
+ Raises:
776
+ NotImplementedError: If runtime doesn't support custom tools
777
+ """
778
+ if not self.supports_custom_tools():
779
+ raise NotImplementedError(
780
+ f"Runtime {self.get_runtime_type().value} does not support custom tools"
781
+ )
782
+ return {
783
+ "format": "unknown",
784
+ "description": "No documentation available",
785
+ "example_code": "",
786
+ "documentation_url": "",
787
+ "required_methods": [],
788
+ "schema": {}
789
+ }
790
+
791
+ def validate_custom_tool(self, tool: Any) -> tuple[bool, Optional[str]]:
792
+ """
793
+ Validate a custom tool for this runtime.
794
+
795
+ Override this method to implement runtime-specific validation logic.
796
+
797
+ Args:
798
+ tool: Custom tool object (format depends on runtime)
799
+
800
+ Returns:
801
+ Tuple of (is_valid, error_message)
802
+ - is_valid: True if tool is valid
803
+ - error_message: Error description if invalid, None if valid
804
+
805
+ Raises:
806
+ NotImplementedError: If runtime doesn't support custom tools
807
+ """
808
+ if not self.supports_custom_tools():
809
+ raise NotImplementedError(
810
+ f"Runtime {self.get_runtime_type().value} does not support custom tools"
811
+ )
812
+ return False, "Validation not implemented"
813
+
814
+ def register_custom_tool(self, tool: Any, metadata: Optional[Dict] = None) -> str:
815
+ """
816
+ Register a custom tool with this runtime.
817
+
818
+ Override this method to implement runtime-specific registration logic.
819
+
820
+ Args:
821
+ tool: Custom tool object (format depends on runtime)
822
+ metadata: Optional metadata (name, description, etc.)
823
+
824
+ Returns:
825
+ Tool identifier for referencing this tool
826
+
827
+ Raises:
828
+ ValueError: If tool is invalid
829
+ NotImplementedError: If runtime doesn't support custom tools
830
+ """
831
+ if not self.supports_custom_tools():
832
+ raise NotImplementedError(
833
+ f"Runtime {self.get_runtime_type().value} does not support custom tools"
834
+ )
835
+
836
+ # Validate first
837
+ is_valid, error = self.validate_custom_tool(tool)
838
+ if not is_valid:
839
+ raise ValueError(f"Invalid custom tool: {error}")
840
+
841
+ raise NotImplementedError("Custom tool registration not implemented")
842
+
843
+ def get_registered_custom_tools(self) -> List[str]:
844
+ """
845
+ Get list of registered custom tool identifiers.
846
+
847
+ Override this method to return the list of tools registered with this runtime.
848
+
849
+ Returns:
850
+ List of tool identifiers
851
+ """
852
+ return []
853
+
854
+
855
+ class RuntimeRegistry:
856
+ """
857
+ Registry for runtime discovery and instantiation.
858
+
859
+ This registry allows runtimes to be:
860
+ - Automatically discovered
861
+ - Registered via decorator
862
+ - Instantiated by name or type
863
+ - Listed for discoverability
864
+ """
865
+
866
+ _registry: Dict[RuntimeType, Type[BaseRuntime]] = {}
867
+
868
+ @classmethod
869
+ def register(cls, runtime_type: RuntimeType):
870
+ """
871
+ Decorator to register a runtime.
872
+
873
+ Usage:
874
+ @RuntimeRegistry.register(RuntimeType.CLAUDE_CODE)
875
+ class ClaudeCodeRuntime(BaseRuntime):
876
+ ...
877
+
878
+ Args:
879
+ runtime_type: Type identifier for this runtime
880
+
881
+ Returns:
882
+ Decorator function
883
+ """
884
+
885
+ def decorator(runtime_class: Type[BaseRuntime]):
886
+ cls._registry[runtime_type] = runtime_class
887
+ logger.info(
888
+ "runtime_registered",
889
+ runtime_type=runtime_type.value,
890
+ runtime_class=runtime_class.__name__,
891
+ )
892
+ return runtime_class
893
+
894
+ return decorator
895
+
896
+ @classmethod
897
+ def get(cls, runtime_type: RuntimeType) -> Type[BaseRuntime]:
898
+ """
899
+ Get runtime class by type.
900
+
901
+ Args:
902
+ runtime_type: Runtime type to get
903
+
904
+ Returns:
905
+ Runtime class
906
+
907
+ Raises:
908
+ ValueError: If runtime type not found
909
+ """
910
+ if runtime_type not in cls._registry:
911
+ raise ValueError(
912
+ f"Runtime type '{runtime_type.value}' not registered. "
913
+ f"Available: {list(cls._registry.keys())}"
914
+ )
915
+ return cls._registry[runtime_type]
916
+
917
+ @classmethod
918
+ def create(
919
+ cls,
920
+ runtime_type: RuntimeType,
921
+ control_plane_client: Any,
922
+ cancellation_manager: Any,
923
+ **kwargs,
924
+ ) -> BaseRuntime:
925
+ """
926
+ Create runtime instance.
927
+
928
+ Args:
929
+ runtime_type: Type of runtime to create
930
+ control_plane_client: Control Plane client
931
+ cancellation_manager: Cancellation manager
932
+ **kwargs: Additional configuration
933
+
934
+ Returns:
935
+ Runtime instance
936
+
937
+ Raises:
938
+ ValueError: If runtime type not found
939
+ """
940
+ runtime_class = cls.get(runtime_type)
941
+ return runtime_class(
942
+ control_plane_client=control_plane_client,
943
+ cancellation_manager=cancellation_manager,
944
+ **kwargs,
945
+ )
946
+
947
+ @classmethod
948
+ def list_available(cls) -> List[RuntimeType]:
949
+ """
950
+ List all registered runtime types.
951
+
952
+ Returns:
953
+ List of available runtime types
954
+ """
955
+ return list(cls._registry.keys())
956
+
957
+ @classmethod
958
+ def get_runtime_info_all(cls) -> Dict[str, Dict[str, Any]]:
959
+ """
960
+ Get information about all registered runtimes.
961
+
962
+ Returns:
963
+ Dict mapping runtime type to info dict
964
+ """
965
+ info = {}
966
+ for runtime_type, runtime_class in cls._registry.items():
967
+ # Instantiate temporarily to get info (mock dependencies)
968
+ try:
969
+ from unittest.mock import MagicMock
970
+
971
+ temp_instance = runtime_class(
972
+ control_plane_client=MagicMock(),
973
+ cancellation_manager=MagicMock(),
974
+ )
975
+ info[runtime_type.value] = temp_instance.get_runtime_info()
976
+ except Exception as e:
977
+ info[runtime_type.value] = {"error": str(e)}
978
+
979
+ return info