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,1585 @@
1
+ """Agent-related Temporal activities"""
2
+
3
+ from dataclasses import dataclass
4
+ from typing import Optional, Any, List, Dict
5
+ from datetime import datetime, timezone
6
+ from temporalio import activity
7
+ from temporalio.exceptions import ApplicationError
8
+ import structlog
9
+ import os
10
+ import httpx
11
+ from pathlib import Path
12
+
13
+ from control_plane_api.worker.utils.logging_helper import execution_logger
14
+
15
+ from agno.tools.shell import ShellTools
16
+ from agno.tools.python import PythonTools
17
+ from agno.tools.file import FileTools
18
+ from control_plane_api.worker.control_plane_client import get_control_plane_client
19
+ from control_plane_api.worker.services.skill_factory import SkillFactory
20
+
21
+ logger = structlog.get_logger()
22
+
23
+ # Global registry for active Agent/Team instances to support cancellation
24
+ # Key: execution_id, Value: {agent: Agent, run_id: str}
25
+ _active_agents: Dict[str, Dict[str, Any]] = {}
26
+
27
+
28
+ def instantiate_skill(skill_data: dict) -> Optional[Any]:
29
+ """
30
+ Instantiate an Agno toolkit based on skill configuration from Control Plane.
31
+
32
+ Args:
33
+ skill_data: Skill data from Control Plane API containing:
34
+ - type: Skill type (file_system, shell, python, docker, etc.)
35
+ - name: Skill name
36
+ - configuration: Dict with skill-specific config
37
+ - enabled: Whether skill is enabled
38
+
39
+ Returns:
40
+ Instantiated Agno toolkit or None if type not supported/enabled
41
+ """
42
+ if not skill_data.get("enabled", True):
43
+ print(f" ⊗ Skipping disabled skill: {skill_data.get('name')}")
44
+ return None
45
+
46
+ skill_type = skill_data.get("type", "").lower()
47
+ config = skill_data.get("configuration", {})
48
+ name = skill_data.get("name", "Unknown")
49
+
50
+ try:
51
+ # Map Control Plane skill types to Agno toolkit classes
52
+ if skill_type in ["file_system", "file", "file_generation"]:
53
+ # FileTools: file operations (read, write, list, search)
54
+ # Note: file_generation is mapped to FileTools (save_file functionality)
55
+ base_dir = config.get("base_dir")
56
+ toolkit = FileTools(
57
+ base_dir=Path(base_dir) if base_dir else None,
58
+ enable_save_file=config.get("enable_save_file", True),
59
+ enable_read_file=config.get("enable_read_file", True),
60
+ enable_list_files=config.get("enable_list_files", True),
61
+ enable_search_files=config.get("enable_search_files", True),
62
+ )
63
+ print(f" ✓ Instantiated FileTools: {name}")
64
+ if skill_type == "file_generation":
65
+ print(f" - Type: File Generation (using FileTools.save_file)")
66
+ print(f" - Base Dir: {base_dir or 'Current directory'}")
67
+ print(f" - Read: {config.get('enable_read_file', True)}, Write: {config.get('enable_save_file', True)}")
68
+ return toolkit
69
+
70
+ elif skill_type in ["shell", "bash"]:
71
+ # ShellTools: shell command execution
72
+ base_dir = config.get("base_dir")
73
+ toolkit = ShellTools(
74
+ base_dir=Path(base_dir) if base_dir else None,
75
+ enable_run_shell_command=config.get("enable_run_shell_command", True),
76
+ )
77
+ print(f" ✓ Instantiated ShellTools: {name}")
78
+ print(f" - Base Dir: {base_dir or 'Current directory'}")
79
+ print(f" - Run Commands: {config.get('enable_run_shell_command', True)}")
80
+ return toolkit
81
+
82
+ elif skill_type == "python":
83
+ # PythonTools: Python code execution
84
+ base_dir = config.get("base_dir")
85
+ toolkit = PythonTools(
86
+ base_dir=Path(base_dir) if base_dir else None,
87
+ safe_globals=config.get("safe_globals"),
88
+ safe_locals=config.get("safe_locals"),
89
+ )
90
+ print(f" ✓ Instantiated PythonTools: {name}")
91
+ print(f" - Base Dir: {base_dir or 'Current directory'}")
92
+ return toolkit
93
+
94
+ elif skill_type == "docker":
95
+ # DockerTools requires docker package and running Docker daemon
96
+ try:
97
+ from agno.tools.docker import DockerTools
98
+ import docker
99
+
100
+ # Check if Docker daemon is accessible
101
+ try:
102
+ docker_client = docker.from_env()
103
+ docker_client.ping()
104
+
105
+ # Docker is available, instantiate toolkit
106
+ toolkit = DockerTools()
107
+ print(f" ✓ Instantiated DockerTools: {name}")
108
+ print(f" - Docker daemon: Connected")
109
+ docker_client.close()
110
+ return toolkit
111
+
112
+ except Exception as docker_error:
113
+ print(f" ⚠ Docker daemon not available - skipping: {name}")
114
+ print(f" Error: {str(docker_error)}")
115
+ return None
116
+
117
+ except ImportError:
118
+ print(f" ⚠ Docker skill requires 'docker' package - skipping: {name}")
119
+ print(f" Install with: pip install docker")
120
+ return None
121
+
122
+ elif skill_type in ["data_visualization", "diagramming", "visualization"]:
123
+ # DataVisualizationTools: Create diagrams using Mermaid syntax
124
+ # This is a custom implementation that uses streaming to send diagram data
125
+ from services.data_visualization import DataVisualizationTools
126
+
127
+ toolkit = DataVisualizationTools(
128
+ max_diagram_size=config.get("max_diagram_size", 50000),
129
+ enable_flowchart=config.get("enable_flowchart", True),
130
+ enable_sequence=config.get("enable_sequence", True),
131
+ enable_class_diagram=config.get("enable_class_diagram", True),
132
+ enable_er_diagram=config.get("enable_er_diagram", True),
133
+ enable_gantt=config.get("enable_gantt", True),
134
+ enable_pie_chart=config.get("enable_pie_chart", True),
135
+ enable_state_diagram=config.get("enable_state_diagram", True),
136
+ enable_git_graph=config.get("enable_git_graph", True),
137
+ enable_user_journey=config.get("enable_user_journey", True),
138
+ enable_quadrant_chart=config.get("enable_quadrant_chart", True),
139
+ )
140
+ print(f" ✓ Instantiated DataVisualizationTools: {name}")
141
+ print(f" - Max diagram size: {config.get('max_diagram_size', 50000)} chars")
142
+ print(f" - Supported: Mermaid diagrams (flowchart, sequence, class, ER, etc.)")
143
+ return toolkit
144
+
145
+ else:
146
+ print(f" ⚠ Unsupported skill type '{skill_type}': {name}")
147
+ return None
148
+
149
+ except Exception as e:
150
+ print(f" ❌ Error instantiating skill '{name}' (type: {skill_type}): {str(e)}")
151
+ logger.error(
152
+ f"Error instantiating skill",
153
+ extra={
154
+ "skill_name": name,
155
+ "skill_type": skill_type,
156
+ "error": str(e)
157
+ }
158
+ )
159
+ return None
160
+
161
+
162
+ @dataclass
163
+ class ActivityExecuteAgentInput:
164
+ """Input for execute_agent_llm activity"""
165
+ execution_id: str
166
+ agent_id: str
167
+ organization_id: str
168
+ prompt: str
169
+ system_prompt: Optional[str] = None
170
+ model_id: Optional[str] = None
171
+ model_config: dict = None
172
+ mcp_servers: dict = None # MCP servers configuration
173
+ session_id: Optional[str] = None # Session ID for Agno session management (use execution_id)
174
+ user_id: Optional[str] = None # User ID for multi-user support
175
+ user_message_id: Optional[str] = None # Message ID from workflow for deduplication
176
+ user_name: Optional[str] = None # User name for attribution
177
+ user_email: Optional[str] = None # User email for attribution
178
+ user_avatar: Optional[str] = None # User avatar for attribution
179
+ # Note: control_plane_url and api_key are read from worker environment variables (CONTROL_PLANE_URL, KUBIYA_API_KEY)
180
+
181
+ def __post_init__(self):
182
+ if self.model_config is None:
183
+ self.model_config = {}
184
+ if self.mcp_servers is None:
185
+ self.mcp_servers = {}
186
+
187
+
188
+ @dataclass
189
+ class ActivityUpdateExecutionInput:
190
+ """Input for update_execution_status activity"""
191
+ execution_id: str
192
+ status: str
193
+ started_at: Optional[str] = None
194
+ completed_at: Optional[str] = None
195
+ response: Optional[str] = None
196
+ error_message: Optional[str] = None
197
+ usage: dict = None
198
+ execution_metadata: dict = None
199
+
200
+ def __post_init__(self):
201
+ if self.usage is None:
202
+ self.usage = {}
203
+ if self.execution_metadata is None:
204
+ self.execution_metadata = {}
205
+
206
+
207
+ @dataclass
208
+ class ActivityGetExecutionInput:
209
+ """Input for get_execution_details activity"""
210
+ execution_id: str
211
+
212
+
213
+ @dataclass
214
+ class ActivityUpdateAgentInput:
215
+ """Input for update_agent_status activity"""
216
+ agent_id: str
217
+ organization_id: str
218
+ status: str
219
+ last_active_at: str
220
+ error_message: Optional[str] = None
221
+ state: dict = None
222
+
223
+ def __post_init__(self):
224
+ if self.state is None:
225
+ self.state = {}
226
+
227
+
228
+ @activity.defn
229
+ async def execute_agent_llm(input: ActivityExecuteAgentInput) -> dict:
230
+ """
231
+ Execute an agent's LLM call with Agno Teams and session management.
232
+
233
+ This activity uses Agno Teams with session support for persistent conversation history.
234
+ The session_id should be set to execution_id for 1:1 mapping.
235
+
236
+ Args:
237
+ input: Activity input with execution details
238
+
239
+ Returns:
240
+ Dict with response, usage, success flag, session messages, etc.
241
+ """
242
+ print("\n" + "="*80)
243
+ print("🤖 AGENT EXECUTION START")
244
+ print("="*80)
245
+ print(f"Execution ID: {input.execution_id}")
246
+ print(f"Agent ID: {input.agent_id}")
247
+ print(f"Organization: {input.organization_id}")
248
+ print(f"Model: {input.model_id or 'default'}")
249
+ print(f"Session ID: {input.session_id}")
250
+ print(f"MCP Servers: {len(input.mcp_servers)} configured" if input.mcp_servers else "MCP Servers: None")
251
+ print(f"Prompt: {input.prompt[:100]}..." if len(input.prompt) > 100 else f"Prompt: {input.prompt}")
252
+ print("="*80 + "\n")
253
+
254
+ activity.logger.info(
255
+ f"Executing agent LLM call with Agno Sessions",
256
+ extra={
257
+ "execution_id": input.execution_id,
258
+ "agent_id": input.agent_id,
259
+ "organization_id": input.organization_id,
260
+ "model_id": input.model_id,
261
+ "has_mcp_servers": bool(input.mcp_servers),
262
+ "mcp_server_count": len(input.mcp_servers) if input.mcp_servers else 0,
263
+ "mcp_server_ids": list(input.mcp_servers.keys()) if input.mcp_servers else [],
264
+ "session_id": input.session_id,
265
+ }
266
+ )
267
+
268
+ try:
269
+ # Get Control Plane client for all communication with Control Plane
270
+ control_plane = get_control_plane_client()
271
+
272
+ # STEP 1: Load existing session history from Control Plane (if this is a continuation)
273
+ # This enables conversation continuity across multiple execution turns
274
+ # IMPORTANT: This must be non-blocking and have proper timeout/retry
275
+ session_history = []
276
+ if input.session_id:
277
+ print(f"\n📥 Loading session history from Control Plane...")
278
+
279
+ # Try up to 3 times with exponential backoff for transient failures
280
+ max_retries = 3
281
+ for attempt in range(max_retries):
282
+ try:
283
+ if attempt > 0:
284
+ print(f" 🔄 Retry attempt {attempt + 1}/{max_retries}...")
285
+
286
+ session_data = control_plane.get_session(
287
+ execution_id=input.execution_id,
288
+ session_id=input.session_id
289
+ )
290
+ if session_data and session_data.get("messages"):
291
+ session_history = session_data["messages"]
292
+ print(f" ✅ Loaded {len(session_history)} messages from previous turns")
293
+
294
+ activity.logger.info(
295
+ "Session history loaded from Control Plane",
296
+ extra={
297
+ "execution_id": input.execution_id,
298
+ "session_id": input.session_id,
299
+ "message_count": len(session_history),
300
+ "attempt": attempt + 1,
301
+ }
302
+ )
303
+ break # Success - exit retry loop
304
+ else:
305
+ print(f" ℹ️ No previous session found - starting new conversation")
306
+ break # No session exists - not an error
307
+
308
+ except httpx.TimeoutException as e:
309
+ print(f" ⏱️ Timeout loading session (attempt {attempt + 1}/{max_retries})")
310
+ activity.logger.warning(
311
+ "Session load timeout",
312
+ extra={"error": str(e), "execution_id": input.execution_id, "attempt": attempt + 1}
313
+ )
314
+ if attempt < max_retries - 1:
315
+ import time
316
+ time.sleep(2 ** attempt) # Exponential backoff: 1s, 2s, 4s
317
+ continue
318
+ else:
319
+ print(f" ⚠️ Session load failed after {max_retries} attempts - continuing without history")
320
+
321
+ except Exception as e:
322
+ error_type = type(e).__name__
323
+ print(f" ⚠️ Failed to load session history ({error_type}): {str(e)[:100]}")
324
+ activity.logger.warning(
325
+ "Failed to load session history from Control Plane",
326
+ extra={
327
+ "error": str(e),
328
+ "error_type": error_type,
329
+ "execution_id": input.execution_id,
330
+ "attempt": attempt + 1
331
+ }
332
+ )
333
+ # For non-timeout errors, don't retry - likely invalid session
334
+ break
335
+
336
+ # Always continue execution even if session loading fails
337
+ print(f" → Continuing with {len(session_history)} messages in context\n")
338
+
339
+ # Get LiteLLM credentials from environment (set by worker from registration)
340
+ litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
341
+ litellm_api_key = os.getenv("LITELLM_API_KEY")
342
+
343
+ if not litellm_api_key:
344
+ raise ValueError("LITELLM_API_KEY environment variable not set")
345
+
346
+ # Get model from input or use default
347
+ model = input.model_id or os.environ.get("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4")
348
+
349
+ # Fetch resolved skills from Control Plane if available
350
+ skills = []
351
+ if input.agent_id:
352
+ print(f"🔧 Fetching skills from Control Plane...")
353
+ try:
354
+ skills = control_plane.get_skills(input.agent_id)
355
+ if skills:
356
+ print(f"✅ Resolved {len(skills)} skills from Control Plane")
357
+ print(f" Skill Types: {[t.get('type') for t in skills]}")
358
+ print(f" Skill Sources: {[t.get('source') for t in skills]}")
359
+ print(f" Skill Names: {[t.get('name') for t in skills]}\n")
360
+
361
+ activity.logger.info(
362
+ f"Resolved skills from Control Plane",
363
+ extra={
364
+ "agent_id": input.agent_id,
365
+ "skill_count": len(skills),
366
+ "skill_types": [t.get("type") for t in skills],
367
+ "skill_sources": [t.get("source") for t in skills],
368
+ "skill_names": [t.get("name") for t in skills],
369
+ }
370
+ )
371
+ else:
372
+ print(f"⚠️ No skills found for agent\n")
373
+ except Exception as e:
374
+ print(f"❌ Error fetching skills: {str(e)}\n")
375
+ activity.logger.error(
376
+ f"Error fetching skills from Control Plane: {str(e)}",
377
+ extra={"error": str(e)}
378
+ )
379
+ # Continue execution without skills
380
+ else:
381
+ print(f"ℹ️ No agent_id provided - skipping skill resolution\n")
382
+
383
+ # Instantiate Agno toolkits from Control Plane skills
384
+ print(f"\n🔧 Instantiating Skills:")
385
+ agno_toolkits = []
386
+ if skills:
387
+ # Create factory instance for agno runtime
388
+ skill_factory = SkillFactory(runtime_type="agno")
389
+ skill_factory.initialize()
390
+
391
+ for skill in skills:
392
+ # Add execution_id to skill data for workflow streaming
393
+ skill['execution_id'] = input.execution_id
394
+ # Use SkillFactory which supports all skill types including workflow_executor
395
+ toolkit = skill_factory.create_skill(skill)
396
+ if toolkit:
397
+ agno_toolkits.append(toolkit)
398
+
399
+ if agno_toolkits:
400
+ print(f"\n✅ Successfully instantiated {len(agno_toolkits)} skill(s)")
401
+ else:
402
+ print(f"\nℹ️ No skills instantiated\n")
403
+
404
+ print(f"📦 Total Tools Available:")
405
+ print(f" MCP Servers: {len(input.mcp_servers)}")
406
+ print(f" OS-Level Skills: {len(agno_toolkits)}\n")
407
+
408
+ activity.logger.info(
409
+ f"Using Agno Agent with sessions and skills",
410
+ extra={
411
+ "execution_id": input.execution_id,
412
+ "session_id": input.session_id,
413
+ "has_mcp_servers": bool(input.mcp_servers),
414
+ "mcp_server_count": len(input.mcp_servers) if input.mcp_servers else 0,
415
+ "mcp_servers": list(input.mcp_servers.keys()) if input.mcp_servers else [],
416
+ "skill_count": len(agno_toolkits),
417
+ "model": model,
418
+ }
419
+ )
420
+
421
+ # Import Agno libraries
422
+ from agno.agent import Agent
423
+ from agno.models.litellm import LiteLLM
424
+
425
+ print(f"\n🤖 Creating Agno Agent:")
426
+ print(f" Model: {model}")
427
+ print(f" Skills: {len(agno_toolkits)}")
428
+
429
+ # Send heartbeat: Creating agent
430
+ activity.heartbeat({"status": "Creating agent with skills..."})
431
+
432
+ # Track tool executions for real-time streaming
433
+ tool_execution_messages = []
434
+
435
+ # Create tool hook to capture tool execution for real-time streaming
436
+ # Agno inspects the signature and passes matching parameters
437
+ def tool_hook(name: str = None, function_name: str = None, function=None, arguments: dict = None, **kwargs):
438
+ """Hook to capture tool execution and add to messages for streaming
439
+
440
+ Agno passes these parameters based on our signature:
441
+ - name or function_name: The tool function name
442
+ - function: The callable being executed (this is the NEXT function in the chain)
443
+ - arguments: Dict of arguments passed to the tool
444
+
445
+ The hook must CALL the function and return its result.
446
+ """
447
+ # Get tool name from Agno's parameters
448
+ tool_name = name or function_name or "unknown"
449
+ tool_args = arguments or {}
450
+
451
+ # Generate unique tool execution ID using UUID to avoid collisions
452
+ import uuid
453
+ tool_execution_id = f"{tool_name}_{uuid.uuid4().hex[:12]}"
454
+
455
+ print(f" 🔧 Tool Starting: {tool_name} (ID: {tool_execution_id})")
456
+ if tool_args:
457
+ args_preview = str(tool_args)[:200]
458
+ print(f" Args: {args_preview}{'...' if len(str(tool_args)) > 200 else ''}")
459
+
460
+ # Publish streaming event to Control Plane (real-time UI update)
461
+ control_plane.publish_event(
462
+ execution_id=input.execution_id,
463
+ event_type="tool_started",
464
+ data={
465
+ "tool_name": tool_name,
466
+ "tool_execution_id": tool_execution_id, # Unique ID for this execution
467
+ "tool_arguments": tool_args,
468
+ "message_id": message_id, # Link tool to assistant message turn
469
+ "message": f"🔧 Executing tool: {tool_name}",
470
+ }
471
+ )
472
+
473
+ tool_execution_messages.append({
474
+ "role": "system",
475
+ "content": f"🔧 Executing tool: **{tool_name}**",
476
+ "tool_name": tool_name,
477
+ "tool_event": "started",
478
+ "timestamp": datetime.now(timezone.utc).isoformat(),
479
+ })
480
+
481
+ # CRITICAL: Actually call the function and handle completion
482
+ result = None
483
+ error = None
484
+ try:
485
+ # Call the actual function (next in the hook chain)
486
+ if function and callable(function):
487
+ result = function(**tool_args) if tool_args else function()
488
+ else:
489
+ raise ValueError(f"Function not callable: {function}")
490
+
491
+ status = "success"
492
+ icon = "✅"
493
+ print(f" {icon} Tool Success: {tool_name}")
494
+
495
+ except Exception as e:
496
+ error = e
497
+ status = "failed"
498
+ icon = "❌"
499
+ print(f" {icon} Tool Failed: {tool_name} - {str(e)}")
500
+
501
+ # Publish completion event to Control Plane (real-time UI update)
502
+ control_plane.publish_event(
503
+ execution_id=input.execution_id,
504
+ event_type="tool_completed",
505
+ data={
506
+ "tool_name": tool_name,
507
+ "tool_execution_id": tool_execution_id, # Same ID to match the started event
508
+ "status": status,
509
+ "error": str(error) if error else None,
510
+ "tool_output": result if result is not None else None, # Include tool output for UI display
511
+ "message_id": message_id, # Link tool to assistant message turn
512
+ "message": f"{icon} Tool {status}: {tool_name}",
513
+ }
514
+ )
515
+
516
+ tool_execution_messages.append({
517
+ "role": "system",
518
+ "content": f"{icon} Tool {status}: **{tool_name}**",
519
+ "tool_name": tool_name,
520
+ "tool_event": "completed",
521
+ "tool_status": status,
522
+ "timestamp": datetime.now(timezone.utc).isoformat(),
523
+ })
524
+
525
+ # If there was an error, re-raise it so Agno knows the tool failed
526
+ if error:
527
+ raise error
528
+
529
+ # Return the result to continue the chain
530
+ return result
531
+
532
+ # Build conversation context from session history for manual session management
533
+ # Workers don't have database access, so we manage sessions via Control Plane API
534
+ conversation_context = []
535
+ if session_history:
536
+ print(f"\n📝 Building conversation context from {len(session_history)} previous messages...")
537
+ for msg in session_history:
538
+ # Convert Control Plane message format to Agno format
539
+ conversation_context.append({
540
+ "role": msg.get("role", "user"),
541
+ "content": msg.get("content", ""),
542
+ })
543
+ print(f" ✅ Conversation context ready\n")
544
+
545
+ # Create Agno Agent with LiteLLM configuration
546
+ # Note: NO database - workers use Control Plane API for session management
547
+ # Use openai/ prefix for custom proxy compatibility
548
+ agent = Agent(
549
+ name=f"Agent {input.agent_id}",
550
+ role=input.system_prompt or "You are a helpful AI assistant",
551
+ model=LiteLLM(
552
+ id=f"openai/{model}",
553
+ api_base=litellm_api_base,
554
+ api_key=litellm_api_key,
555
+ ),
556
+ tools=agno_toolkits if agno_toolkits else None, # Add skills to agent
557
+ tool_hooks=[tool_hook], # Add hook for real-time tool updates
558
+ # NO db parameter - session management via Control Plane API
559
+ )
560
+
561
+ # Register agent for cancellation support
562
+ _active_agents[input.execution_id] = {
563
+ "agent": agent,
564
+ "run_id": None, # Will be set when run starts
565
+ "started_at": datetime.now(timezone.utc).isoformat(),
566
+ }
567
+ print(f"✅ Agent registered for cancellation support (execution_id: {input.execution_id})\n")
568
+
569
+ # Cache execution metadata in Redis for fast SSE lookups (avoid DB queries)
570
+ control_plane.cache_metadata(input.execution_id, "AGENT")
571
+
572
+ # Execute agent run with streaming
573
+ print("⚡ Executing Agent Run with Streaming...\n")
574
+
575
+ # Send heartbeat: Starting execution
576
+ activity.heartbeat({"status": "Agent is processing your request..."})
577
+
578
+ import asyncio
579
+
580
+ # Stream the response and collect chunks
581
+ response_chunks = []
582
+ full_response = ""
583
+
584
+ # Generate unique message ID for this turn (execution_id + timestamp)
585
+ import time
586
+ message_id = f"{input.execution_id}_{int(time.time() * 1000000)}"
587
+
588
+ def stream_agent_run():
589
+ """Run agent with streaming and collect response"""
590
+ nonlocal full_response
591
+ run_id_published = False
592
+
593
+ try:
594
+ # Build full prompt with conversation history for context
595
+ # Since worker has no database, we manually prepend history
596
+ if conversation_context:
597
+ # Agno Agent supports passing messages parameter for conversation context
598
+ run_response = agent.run(
599
+ input.prompt,
600
+ stream=True,
601
+ messages=conversation_context, # Pass previous conversation
602
+ )
603
+ else:
604
+ # First turn - no history
605
+ run_response = agent.run(input.prompt, stream=True)
606
+
607
+ # Iterate over streaming chunks
608
+ for chunk in run_response:
609
+ # Capture and publish run_id from first chunk for cancellation support
610
+ if not run_id_published and hasattr(chunk, 'run_id') and chunk.run_id:
611
+ agno_run_id = chunk.run_id
612
+ print(f"\n🆔 Agno run_id: {agno_run_id}")
613
+
614
+ # Store run_id in registry for cancellation
615
+ if input.execution_id in _active_agents:
616
+ _active_agents[input.execution_id]["run_id"] = agno_run_id
617
+
618
+ # Publish run_id to Redis for Control Plane cancellation access
619
+ # This allows users to cancel via STOP button in UI
620
+ control_plane.publish_event(
621
+ execution_id=input.execution_id,
622
+ event_type="run_started",
623
+ data={
624
+ "run_id": agno_run_id,
625
+ "agent_id": input.agent_id,
626
+ "cancellable": True,
627
+ }
628
+ )
629
+ run_id_published = True
630
+
631
+ # Filter out whitespace-only chunks to prevent "(no content)" in UI
632
+ if hasattr(chunk, 'content') and chunk.content:
633
+ content = str(chunk.content)
634
+ # Only process and send chunks with non-whitespace content
635
+ if content.strip():
636
+ full_response += content
637
+ response_chunks.append(content)
638
+ print(content, end='', flush=True)
639
+
640
+ # Stream chunk to Control Plane for real-time UI updates
641
+ # Include message_id so UI knows which message these chunks belong to
642
+ control_plane.publish_event(
643
+ execution_id=input.execution_id,
644
+ event_type="message_chunk",
645
+ data={
646
+ "role": "assistant",
647
+ "content": content,
648
+ "is_chunk": True,
649
+ "message_id": message_id, # Unique ID for this turn
650
+ }
651
+ )
652
+
653
+ # Note: Cannot send heartbeat from sync context (thread pool)
654
+
655
+ print() # New line after streaming
656
+
657
+ # Return the iterator's final result
658
+ return run_response
659
+ except Exception as e:
660
+ print(f"\n❌ Streaming error: {str(e)}")
661
+ # Fall back to non-streaming
662
+ if conversation_context:
663
+ return agent.run(input.prompt, stream=False, messages=conversation_context)
664
+ else:
665
+ return agent.run(input.prompt, stream=False)
666
+
667
+ # Execute in thread pool (NO TIMEOUT - tasks can run as long as needed)
668
+ # Control Plane can cancel via Agno's cancel_run API if user requests it
669
+ result = await asyncio.to_thread(stream_agent_run)
670
+
671
+ # Send heartbeat: Completed
672
+ activity.heartbeat({"status": "Agent execution completed, preparing response..."})
673
+
674
+ print("✅ Agent Execution Completed!")
675
+ print(f" Response Length: {len(full_response)} chars\n")
676
+
677
+ activity.logger.info(
678
+ f"Agent LLM call completed",
679
+ extra={
680
+ "execution_id": input.execution_id,
681
+ "has_content": bool(full_response),
682
+ }
683
+ )
684
+
685
+ # Use the streamed response content
686
+ response_content = full_response if full_response else (result.content if hasattr(result, "content") else str(result))
687
+
688
+ # Extract tool call messages for UI streaming
689
+ tool_messages = []
690
+ if hasattr(result, "messages") and result.messages:
691
+ for msg in result.messages:
692
+ # Check if message has tool calls
693
+ if hasattr(msg, "tool_calls") and msg.tool_calls:
694
+ for tool_call in msg.tool_calls:
695
+ tool_name = getattr(tool_call, "function", {}).get("name") if hasattr(tool_call, "function") else str(tool_call)
696
+ tool_args = getattr(tool_call, "function", {}).get("arguments") if hasattr(tool_call, "function") else {}
697
+
698
+ print(f" 🔧 Tool Call: {tool_name}")
699
+
700
+ tool_messages.append({
701
+ "role": "tool",
702
+ "content": f"Executing {tool_name}...",
703
+ "tool_name": tool_name,
704
+ "tool_execution_id": tool_execution_id, # CRITICAL: For deduplication
705
+ "tool_input": tool_args,
706
+ "timestamp": datetime.now(timezone.utc).isoformat(),
707
+ })
708
+
709
+ if tool_messages:
710
+ print(f"\n🔧 Tool Calls Captured: {len(tool_messages)}")
711
+
712
+ # Extract usage metrics if available
713
+ usage = {}
714
+ if hasattr(result, "metrics") and result.metrics:
715
+ metrics = result.metrics
716
+ usage = {
717
+ "prompt_tokens": getattr(metrics, "input_tokens", 0),
718
+ "completion_tokens": getattr(metrics, "output_tokens", 0),
719
+ "total_tokens": getattr(metrics, "total_tokens", 0),
720
+ }
721
+ print(f"📊 Token Usage:")
722
+ print(f" Input Tokens: {usage.get('prompt_tokens', 0)}")
723
+ print(f" Output Tokens: {usage.get('completion_tokens', 0)}")
724
+ print(f" Total Tokens: {usage.get('total_tokens', 0)}\n")
725
+
726
+ print(f"📝 Response Preview:")
727
+ print(f" {response_content[:200]}..." if len(response_content) > 200 else f" {response_content}")
728
+
729
+ # CRITICAL: Persist COMPLETE session history to Control Plane API
730
+ # This includes previous history + current turn for conversation continuity
731
+ # IMPORTANT: Use retry logic - persistence failures shouldn't break execution
732
+ print("\n💾 Persisting session history to Control Plane...")
733
+
734
+ # Build complete session: previous history + current turn's messages
735
+ updated_session_messages = list(session_history) # Start with loaded history
736
+
737
+ # Track proper timestamps for chronological ordering
738
+ # User message happened at execution start, assistant response completed now
739
+ from datetime import timedelta
740
+ response_completed_at = datetime.now(timezone.utc)
741
+ # User message must be BEFORE assistant message (subtract execution time)
742
+ # Use a conservative estimate: user message was at least 1 second before response
743
+ execution_started_at = response_completed_at - timedelta(seconds=1)
744
+
745
+ # Add current turn messages (user prompt + assistant response)
746
+ # CRITICAL: User message must have EARLIER timestamp than assistant for proper ordering
747
+ # CRITICAL: Use message_id from workflow signal if available (for deduplication consistency)
748
+ user_message_id = getattr(input, "user_message_id", None) or f"{input.execution_id}_user_{int(execution_started_at.timestamp() * 1000000)}"
749
+
750
+ current_turn_messages = [
751
+ {
752
+ "role": "user",
753
+ "content": input.prompt,
754
+ "timestamp": execution_started_at.isoformat(), # When user sent message
755
+ "user_id": input.user_id,
756
+ "user_name": getattr(input, "user_name", None),
757
+ "user_email": getattr(input, "user_email", None),
758
+ "user_avatar": getattr(input, "user_avatar", None),
759
+ "message_id": user_message_id, # Use provided message_id or generate stable one
760
+ },
761
+ {
762
+ "role": "assistant",
763
+ "content": response_content,
764
+ "timestamp": response_completed_at.isoformat(), # When assistant finished responding
765
+ "message_id": message_id, # Use the message_id generated at start of execution
766
+ }
767
+ ]
768
+
769
+ print(f" 📝 Adding {len(current_turn_messages)} messages from current turn (user + assistant)...")
770
+ updated_session_messages.extend(current_turn_messages)
771
+
772
+ # CRITICAL: Also add tool messages to session history for proper rendering
773
+ if tool_messages:
774
+ print(f" 🔧 Adding {len(tool_messages)} tool messages to session history...")
775
+ updated_session_messages.extend(tool_messages)
776
+
777
+ if updated_session_messages:
778
+ # Try up to 3 times to persist session history
779
+ max_retries = 3
780
+ persisted = False
781
+
782
+ for attempt in range(max_retries):
783
+ try:
784
+ if attempt > 0:
785
+ print(f" 🔄 Retry persistence attempt {attempt + 1}/{max_retries}...")
786
+
787
+ success = control_plane.persist_session(
788
+ execution_id=input.execution_id,
789
+ session_id=input.session_id or input.execution_id,
790
+ user_id=input.user_id,
791
+ messages=updated_session_messages, # Complete conversation history
792
+ metadata={
793
+ "agent_id": input.agent_id,
794
+ "organization_id": input.organization_id,
795
+ "turn_count": len(updated_session_messages),
796
+ }
797
+ )
798
+
799
+ if success:
800
+ print(f" ✅ Complete session history persisted ({len(updated_session_messages)} total messages)")
801
+ persisted = True
802
+ break
803
+ else:
804
+ print(f" ⚠️ Persistence failed (attempt {attempt + 1}/{max_retries})")
805
+ if attempt < max_retries - 1:
806
+ import time
807
+ time.sleep(2 ** attempt) # Exponential backoff
808
+
809
+ except Exception as session_error:
810
+ error_type = type(session_error).__name__
811
+ print(f" ⚠️ Persistence error ({error_type}, attempt {attempt + 1}/{max_retries})")
812
+ logger.warning(
813
+ "session_persistence_error",
814
+ extra={
815
+ "error": str(session_error),
816
+ "error_type": error_type,
817
+ "execution_id": input.execution_id,
818
+ "attempt": attempt + 1
819
+ }
820
+ )
821
+ if attempt < max_retries - 1:
822
+ import time
823
+ time.sleep(2 ** attempt) # Exponential backoff
824
+
825
+ if not persisted:
826
+ print(f" ⚠️ Session persistence failed after {max_retries} attempts")
827
+ logger.error(
828
+ "session_persistence_failed_all_retries",
829
+ extra={
830
+ "execution_id": input.execution_id,
831
+ "message_count": len(updated_session_messages)
832
+ }
833
+ )
834
+ # Don't fail execution - session loss is better than execution failure
835
+ else:
836
+ print(" ℹ️ No messages - skipping session persistence")
837
+
838
+ print("\n" + "="*80)
839
+ print("🏁 AGENT EXECUTION END")
840
+ print("="*80 + "\n")
841
+
842
+ # Cleanup: Remove agent from registry
843
+ if input.execution_id in _active_agents:
844
+ del _active_agents[input.execution_id]
845
+ print(f"✅ Agent unregistered (execution_id: {input.execution_id})\n")
846
+
847
+ return {
848
+ "success": True,
849
+ "response": response_content,
850
+ "usage": usage,
851
+ "model": model,
852
+ "finish_reason": "stop",
853
+ "mcp_tools_used": 0, # TODO: Track MCP tool usage
854
+ "tool_messages": tool_messages, # Include tool call messages for UI
855
+ "tool_execution_messages": tool_execution_messages, # Include real-time tool execution status
856
+ }
857
+
858
+ except Exception as e:
859
+ # Cleanup on error
860
+ if input.execution_id in _active_agents:
861
+ del _active_agents[input.execution_id]
862
+
863
+ # Ensure error message is never empty
864
+ error_msg = str(e) or repr(e) or f"{type(e).__name__}: No error details available"
865
+
866
+ print("\n" + "="*80)
867
+ print("❌ AGENT EXECUTION FAILED")
868
+ print("="*80)
869
+ print(f"Error: {error_msg}")
870
+ print(f"Error Type: {type(e).__name__}")
871
+ print("="*80 + "\n")
872
+
873
+ activity.logger.error(
874
+ f"Agent LLM call failed",
875
+ extra={
876
+ "execution_id": input.execution_id,
877
+ "error": error_msg,
878
+ "error_type": type(e).__name__,
879
+ },
880
+ exc_info=True,
881
+ )
882
+
883
+ # Raise ApplicationError so Temporal marks the workflow as FAILED
884
+ raise ApplicationError(
885
+ f"Agent execution failed: {error_msg}",
886
+ non_retryable=False, # Allow retries per retry policy
887
+ type=type(e).__name__
888
+ )
889
+
890
+
891
+ @activity.defn
892
+ async def update_execution_status(input: ActivityUpdateExecutionInput) -> dict:
893
+ """
894
+ Update execution status in database via Control Plane API.
895
+
896
+ This activity calls the Control Plane API to update execution records.
897
+ Also records which worker processed this execution.
898
+
899
+ Args:
900
+ input: Activity input with update details
901
+
902
+ Returns:
903
+ Dict with success flag
904
+ """
905
+ execution_logger.activity_started(
906
+ "Update Status",
907
+ input.execution_id,
908
+ details={"status": input.status}
909
+ )
910
+
911
+ try:
912
+ # Get Control Plane URL and Kubiya API key from environment
913
+ control_plane_url = os.getenv("CONTROL_PLANE_URL")
914
+ kubiya_api_key = os.getenv("KUBIYA_API_KEY")
915
+ worker_id = os.getenv("WORKER_ID", "unknown")
916
+
917
+ if not control_plane_url:
918
+ raise ValueError("CONTROL_PLANE_URL environment variable not set")
919
+ if not kubiya_api_key:
920
+ raise ValueError("KUBIYA_API_KEY environment variable not set")
921
+
922
+ # Collect worker system information
923
+ import socket
924
+ import platform
925
+ worker_info = {
926
+ "worker_id": worker_id,
927
+ "hostname": socket.gethostname(),
928
+ "platform": platform.platform(),
929
+ "python_version": platform.python_version(),
930
+ }
931
+
932
+ # Build update payload
933
+ update_payload = {}
934
+
935
+ if input.status:
936
+ update_payload["status"] = input.status
937
+
938
+ if input.started_at:
939
+ update_payload["started_at"] = input.started_at
940
+
941
+ if input.completed_at:
942
+ update_payload["completed_at"] = input.completed_at
943
+
944
+ if input.response is not None:
945
+ update_payload["response"] = input.response
946
+
947
+ if input.error_message is not None:
948
+ update_payload["error_message"] = input.error_message
949
+
950
+ if input.usage:
951
+ update_payload["usage"] = input.usage
952
+
953
+ # Merge worker info into execution_metadata
954
+ execution_metadata = input.execution_metadata or {}
955
+ if not execution_metadata.get("worker_info"):
956
+ execution_metadata["worker_info"] = worker_info
957
+ update_payload["execution_metadata"] = execution_metadata
958
+
959
+ # Call Control Plane API
960
+ async with httpx.AsyncClient(timeout=30.0) as client:
961
+ response = await client.patch(
962
+ f"{control_plane_url}/api/v1/executions/{input.execution_id}",
963
+ json=update_payload,
964
+ headers={
965
+ "Authorization": f"Bearer {kubiya_api_key}",
966
+ "Content-Type": "application/json",
967
+ }
968
+ )
969
+
970
+ if response.status_code == 404:
971
+ # Execution not found - this is not a retryable error
972
+ # The execution may have been deleted, or never existed
973
+ execution_logger.activity_failed(
974
+ "Update Status",
975
+ input.execution_id,
976
+ "Execution record not found - may have been deleted or cancelled",
977
+ will_retry=False
978
+ )
979
+
980
+ # Raise ApplicationError with non_retryable=True since execution doesn't exist
981
+ raise ApplicationError(
982
+ f"Execution {input.execution_id} not found - may have been deleted",
983
+ non_retryable=True,
984
+ type="ExecutionNotFound"
985
+ )
986
+ elif response.status_code != 200:
987
+ raise Exception(f"Failed to update execution: {response.status_code} - {response.text}")
988
+
989
+ # IMPORTANT: Log status changes to RUNNING prominently
990
+ if input.status == "running":
991
+ logger.info(
992
+ "execution_status_updated_to_running",
993
+ execution_id=input.execution_id[:8] if input.execution_id else "unknown",
994
+ status=input.status,
995
+ worker_id=worker_id
996
+ )
997
+
998
+ execution_logger.activity_completed(
999
+ "Update Status",
1000
+ input.execution_id,
1001
+ result=f"Status: {input.status}"
1002
+ )
1003
+
1004
+ activity.logger.info(
1005
+ f"Execution status updated via API",
1006
+ extra={
1007
+ "execution_id": input.execution_id,
1008
+ "status": input.status,
1009
+ }
1010
+ )
1011
+
1012
+ return {"success": True, "execution_not_found": False}
1013
+
1014
+ except Exception as e:
1015
+ print(f"❌ Failed to update status: {str(e)}\n")
1016
+
1017
+ activity.logger.error(
1018
+ f"Failed to update execution status",
1019
+ extra={
1020
+ "execution_id": input.execution_id,
1021
+ "error": str(e),
1022
+ }
1023
+ )
1024
+ raise
1025
+
1026
+
1027
+ @activity.defn
1028
+ async def get_execution_details(input: ActivityGetExecutionInput) -> dict:
1029
+ """
1030
+ Get execution details from Control Plane API.
1031
+
1032
+ This activity fetches the current execution state including status and metadata.
1033
+
1034
+ Args:
1035
+ input: Activity input with execution_id
1036
+
1037
+ Returns:
1038
+ Dict with execution details including status
1039
+ """
1040
+ execution_logger.activity_started(
1041
+ "Get Execution Details",
1042
+ input.execution_id
1043
+ )
1044
+
1045
+ try:
1046
+ # Get Control Plane URL and Kubiya API key from environment
1047
+ control_plane_url = os.getenv("CONTROL_PLANE_URL")
1048
+ kubiya_api_key = os.getenv("KUBIYA_API_KEY")
1049
+
1050
+ if not control_plane_url:
1051
+ raise ValueError("CONTROL_PLANE_URL environment variable not set")
1052
+ if not kubiya_api_key:
1053
+ raise ValueError("KUBIYA_API_KEY environment variable not set")
1054
+
1055
+ # Call Control Plane API to get execution
1056
+ async with httpx.AsyncClient(timeout=30.0) as client:
1057
+ response = await client.get(
1058
+ f"{control_plane_url}/api/v1/executions/{input.execution_id}",
1059
+ headers={
1060
+ "Authorization": f"Bearer {kubiya_api_key}",
1061
+ "Content-Type": "application/json",
1062
+ }
1063
+ )
1064
+
1065
+ if response.status_code == 404:
1066
+ execution_logger.activity_failed(
1067
+ "Get Execution Details",
1068
+ input.execution_id,
1069
+ "Execution record not found",
1070
+ will_retry=False
1071
+ )
1072
+
1073
+ raise ApplicationError(
1074
+ f"Execution {input.execution_id} not found",
1075
+ non_retryable=True,
1076
+ type="ExecutionNotFound"
1077
+ )
1078
+ elif response.status_code != 200:
1079
+ raise Exception(f"Failed to get execution: {response.status_code} - {response.text}")
1080
+
1081
+ execution_data = response.json()
1082
+
1083
+ execution_logger.activity_completed(
1084
+ "Get Execution Details",
1085
+ input.execution_id,
1086
+ result=f"Status: {execution_data.get('status', 'unknown')}"
1087
+ )
1088
+
1089
+ activity.logger.info(
1090
+ f"Execution details fetched via API",
1091
+ extra={
1092
+ "execution_id": input.execution_id,
1093
+ "status": execution_data.get("status"),
1094
+ }
1095
+ )
1096
+
1097
+ return execution_data
1098
+
1099
+ except Exception as e:
1100
+ activity.logger.error(
1101
+ f"Failed to get execution details",
1102
+ extra={
1103
+ "execution_id": input.execution_id,
1104
+ "error": str(e),
1105
+ }
1106
+ )
1107
+ raise
1108
+
1109
+
1110
+ @activity.defn
1111
+ async def update_agent_status(input: ActivityUpdateAgentInput) -> dict:
1112
+ """
1113
+ Update agent status in database via Control Plane API.
1114
+
1115
+ This activity calls the Control Plane API to update agent records.
1116
+
1117
+ Args:
1118
+ input: Activity input with update details
1119
+
1120
+ Returns:
1121
+ Dict with success flag
1122
+ """
1123
+ activity.logger.info(
1124
+ f"Updating agent status via Control Plane API",
1125
+ extra={
1126
+ "agent_id": input.agent_id,
1127
+ "status": input.status,
1128
+ }
1129
+ )
1130
+
1131
+ try:
1132
+ # Get Control Plane URL and Kubiya API key from environment
1133
+ control_plane_url = os.getenv("CONTROL_PLANE_URL")
1134
+ kubiya_api_key = os.getenv("KUBIYA_API_KEY")
1135
+
1136
+ if not control_plane_url:
1137
+ raise ValueError("CONTROL_PLANE_URL environment variable not set")
1138
+ if not kubiya_api_key:
1139
+ raise ValueError("KUBIYA_API_KEY environment variable not set")
1140
+
1141
+ # Build update payload
1142
+ update_payload = {
1143
+ "status": input.status,
1144
+ "last_active_at": input.last_active_at,
1145
+ }
1146
+
1147
+ if input.error_message is not None:
1148
+ update_payload["error_message"] = input.error_message
1149
+
1150
+ if input.state:
1151
+ update_payload["state"] = input.state
1152
+
1153
+ # Call Control Plane API
1154
+ async with httpx.AsyncClient(timeout=30.0) as client:
1155
+ response = await client.patch(
1156
+ f"{control_plane_url}/api/v1/agents/{input.agent_id}",
1157
+ json=update_payload,
1158
+ headers={
1159
+ "Authorization": f"Bearer {kubiya_api_key}",
1160
+ "Content-Type": "application/json",
1161
+ }
1162
+ )
1163
+
1164
+ # For team executions, the "agent_id" is actually a team_id, so it won't be found in agents table
1165
+ # This is expected and not an error - just log and return success
1166
+ if response.status_code == 404:
1167
+ activity.logger.info(
1168
+ f"Agent not found (likely a team execution) - skipping agent status update",
1169
+ extra={
1170
+ "agent_id": input.agent_id,
1171
+ "status": input.status,
1172
+ }
1173
+ )
1174
+ return {"success": True, "skipped": True}
1175
+ elif response.status_code != 200:
1176
+ raise Exception(f"Failed to update agent: {response.status_code} - {response.text}")
1177
+
1178
+ activity.logger.info(
1179
+ f"Agent status updated via API",
1180
+ extra={
1181
+ "agent_id": input.agent_id,
1182
+ "status": input.status,
1183
+ }
1184
+ )
1185
+
1186
+ return {"success": True}
1187
+
1188
+ except Exception as e:
1189
+ activity.logger.error(
1190
+ f"Failed to update agent status",
1191
+ extra={
1192
+ "agent_id": input.agent_id,
1193
+ "error": str(e),
1194
+ }
1195
+ )
1196
+ raise
1197
+
1198
+
1199
+ @dataclass
1200
+ class ActivityCancelAgentInput:
1201
+ execution_id: str
1202
+
1203
+
1204
+ @activity.defn(name="cancel_agent_run")
1205
+ async def cancel_agent_run(input: ActivityCancelAgentInput) -> dict:
1206
+ """
1207
+ Cancel an active agent/team run using Agno's cancel_run API.
1208
+
1209
+ This is called when a user clicks the STOP button in the UI.
1210
+ Uses the global registry to find the Agent instance and run_id,
1211
+ then calls agent.cancel_run(run_id) to stop execution.
1212
+
1213
+ Args:
1214
+ input: Contains execution_id to identify which run to cancel
1215
+
1216
+ Returns:
1217
+ Dict with success status and cancellation details
1218
+ """
1219
+ print("\n" + "="*80)
1220
+ print("🛑 CANCEL AGENT RUN")
1221
+ print("="*80)
1222
+ print(f"Execution ID: {input.execution_id}\n")
1223
+
1224
+ try:
1225
+ # Look up agent in registry
1226
+ if input.execution_id not in _active_agents:
1227
+ print(f"⚠️ Agent not found in registry - may have already completed")
1228
+ activity.logger.warning(
1229
+ "cancel_agent_not_found",
1230
+ extra={"execution_id": input.execution_id}
1231
+ )
1232
+ return {
1233
+ "success": False,
1234
+ "error": "Agent not found or already completed",
1235
+ "execution_id": input.execution_id,
1236
+ }
1237
+
1238
+ agent_info = _active_agents[input.execution_id]
1239
+ agent = agent_info["agent"]
1240
+ run_id = agent_info.get("run_id")
1241
+
1242
+ if not run_id:
1243
+ print(f"⚠️ No run_id found - execution may not have started yet")
1244
+ return {
1245
+ "success": False,
1246
+ "error": "Execution not started yet",
1247
+ "execution_id": input.execution_id,
1248
+ }
1249
+
1250
+ print(f"🆔 Found run_id: {run_id}")
1251
+ print(f"🛑 Calling agent.cancel_run()...")
1252
+
1253
+ # Call Agno's cancel_run API
1254
+ success = agent.cancel_run(run_id)
1255
+
1256
+ if success:
1257
+ print(f"✅ Agent run cancelled successfully!\n")
1258
+ activity.logger.info(
1259
+ "agent_run_cancelled",
1260
+ extra={
1261
+ "execution_id": input.execution_id,
1262
+ "run_id": run_id,
1263
+ }
1264
+ )
1265
+
1266
+ # Clean up registry
1267
+ del _active_agents[input.execution_id]
1268
+
1269
+ return {
1270
+ "success": True,
1271
+ "execution_id": input.execution_id,
1272
+ "run_id": run_id,
1273
+ "cancelled_at": datetime.now(timezone.utc).isoformat(),
1274
+ }
1275
+ else:
1276
+ print(f"⚠️ Cancel failed - run may have already completed\n")
1277
+ return {
1278
+ "success": False,
1279
+ "error": "Cancel failed - run may be completed",
1280
+ "execution_id": input.execution_id,
1281
+ "run_id": run_id,
1282
+ }
1283
+
1284
+ except Exception as e:
1285
+ print(f"❌ Error cancelling run: {str(e)}\n")
1286
+ activity.logger.error(
1287
+ "cancel_agent_error",
1288
+ extra={
1289
+ "execution_id": input.execution_id,
1290
+ "error": str(e),
1291
+ }
1292
+ )
1293
+ return {
1294
+ "success": False,
1295
+ "error": str(e),
1296
+ "execution_id": input.execution_id,
1297
+ }
1298
+
1299
+
1300
+ @dataclass
1301
+ class ActivityPersistConversationInput:
1302
+ """Input for persisting conversation history"""
1303
+ execution_id: str
1304
+ session_id: str
1305
+ messages: List[Dict[str, Any]]
1306
+ user_id: Optional[str] = None
1307
+ metadata: Optional[Dict[str, Any]] = None
1308
+
1309
+
1310
+ @activity.defn(name="persist_conversation_history")
1311
+ async def persist_conversation_history(input: ActivityPersistConversationInput) -> dict:
1312
+ """
1313
+ Persist conversation history to Control Plane after each turn.
1314
+
1315
+ This ensures conversation state is saved end-to-end, making it available:
1316
+ - For future turns in the same conversation
1317
+ - For UI display and history views
1318
+ - For analytics and debugging
1319
+ - Even if the worker crashes or restarts
1320
+
1321
+ The Control Plane stores this in the database and caches it in Redis
1322
+ for fast retrieval on subsequent turns.
1323
+
1324
+ Args:
1325
+ input: Contains execution_id, session_id, messages, and optional metadata
1326
+
1327
+ Returns:
1328
+ Dict with success status and persistence details
1329
+ """
1330
+ execution_id_short = input.execution_id[:8] if input.execution_id else "unknown"
1331
+
1332
+ activity.logger.info(
1333
+ "persisting_conversation",
1334
+ extra={
1335
+ "execution_id": execution_id_short,
1336
+ "session_id": input.session_id[:8] if input.session_id else "none",
1337
+ "message_count": len(input.messages),
1338
+ }
1339
+ )
1340
+
1341
+ try:
1342
+ # Get Control Plane client
1343
+ control_plane = get_control_plane_client()
1344
+
1345
+ # CRITICAL: Deduplicate messages by message_id before persisting
1346
+ # This prevents duplicate messages from being stored in the database
1347
+ seen_message_ids = set()
1348
+ deduplicated_messages = []
1349
+
1350
+ for msg in input.messages:
1351
+ # CRITICAL: For tool messages, use tool_execution_id as unique identifier
1352
+ # This prevents duplicate tool calls from being persisted
1353
+ if msg.get("role") == "tool" and msg.get("tool_execution_id"):
1354
+ msg_id = msg["tool_execution_id"]
1355
+ activity.logger.debug(
1356
+ "using_tool_execution_id_for_deduplication",
1357
+ extra={
1358
+ "execution_id": execution_id_short,
1359
+ "tool_execution_id": msg_id,
1360
+ "tool_name": msg.get("tool_name")
1361
+ }
1362
+ )
1363
+ else:
1364
+ msg_id = msg.get("message_id")
1365
+
1366
+ if not msg_id:
1367
+ # Generate stable message_id for messages without one
1368
+ # Use timestamp + role to create consistent ID
1369
+ timestamp_str = msg.get("timestamp", "")
1370
+ role = msg.get("role", "unknown")
1371
+ msg_id = f"{input.execution_id}_{role}_{timestamp_str}"
1372
+ msg["message_id"] = msg_id
1373
+ activity.logger.debug(
1374
+ "generated_missing_message_id",
1375
+ extra={
1376
+ "execution_id": execution_id_short,
1377
+ "role": role,
1378
+ "generated_id": msg_id
1379
+ }
1380
+ )
1381
+
1382
+ # Check for duplicates
1383
+ if msg_id in seen_message_ids:
1384
+ activity.logger.warning(
1385
+ "skipping_duplicate_message_id",
1386
+ extra={
1387
+ "execution_id": execution_id_short,
1388
+ "message_id": msg_id,
1389
+ "role": msg.get("role"),
1390
+ "content_preview": msg.get("content", "")[:50]
1391
+ }
1392
+ )
1393
+ continue
1394
+
1395
+ seen_message_ids.add(msg_id)
1396
+ deduplicated_messages.append(msg)
1397
+
1398
+ activity.logger.info(
1399
+ "deduplication_complete",
1400
+ extra={
1401
+ "execution_id": execution_id_short,
1402
+ "before": len(input.messages),
1403
+ "after": len(deduplicated_messages),
1404
+ "removed": len(input.messages) - len(deduplicated_messages)
1405
+ }
1406
+ )
1407
+
1408
+ # Persist conversation via Control Plane API with deduplicated messages
1409
+ success = control_plane.persist_session(
1410
+ execution_id=input.execution_id,
1411
+ session_id=input.session_id or input.execution_id,
1412
+ user_id=input.user_id,
1413
+ messages=deduplicated_messages, # Use deduplicated messages
1414
+ metadata=input.metadata or {}
1415
+ )
1416
+
1417
+ if success:
1418
+ activity.logger.info(
1419
+ "conversation_persisted",
1420
+ extra={
1421
+ "execution_id": execution_id_short,
1422
+ "message_count": len(input.messages),
1423
+ }
1424
+ )
1425
+ return {
1426
+ "success": True,
1427
+ "execution_id": input.execution_id,
1428
+ "message_count": len(input.messages),
1429
+ "persisted_at": datetime.now(timezone.utc).isoformat(),
1430
+ }
1431
+ else:
1432
+ activity.logger.warning(
1433
+ "conversation_persistence_failed",
1434
+ extra={
1435
+ "execution_id": execution_id_short,
1436
+ }
1437
+ )
1438
+ return {
1439
+ "success": False,
1440
+ "error": "Control Plane API returned failure",
1441
+ "execution_id": input.execution_id,
1442
+ }
1443
+
1444
+ except Exception as e:
1445
+ error_type = type(e).__name__
1446
+ error_msg = str(e) if str(e) else "No error message provided"
1447
+
1448
+ activity.logger.error(
1449
+ "conversation_persistence_error",
1450
+ extra={
1451
+ "execution_id": execution_id_short,
1452
+ "error_type": error_type,
1453
+ "error": error_msg[:500], # Truncate very long errors
1454
+ "message_count": len(input.messages),
1455
+ },
1456
+ exc_info=True,
1457
+ )
1458
+ return {
1459
+ "success": False,
1460
+ "error": f"{error_type}: {error_msg}",
1461
+ "error_type": error_type,
1462
+ "execution_id": input.execution_id,
1463
+ }
1464
+
1465
+
1466
+ @dataclass
1467
+ class AnalyticsActivityInput:
1468
+ """Input for analytics submission activity"""
1469
+ execution_id: str
1470
+ turn_number: int
1471
+ result: Dict[str, Any] # RuntimeExecutionResult as dict
1472
+ turn_start_time: float
1473
+
1474
+
1475
+ @activity.defn
1476
+ async def submit_runtime_analytics_activity(input: AnalyticsActivityInput) -> dict:
1477
+ """
1478
+ Temporal activity for submitting runtime analytics.
1479
+
1480
+ Runs independently from main execution flow, with its own timeout and retry logic.
1481
+ Failures are logged but do not affect execution success.
1482
+ """
1483
+ from control_plane_api.worker.services.runtime_analytics import submit_runtime_analytics
1484
+ from control_plane_api.worker.services.analytics_service import AnalyticsService
1485
+ from control_plane_api.worker.runtimes.base import RuntimeExecutionResult
1486
+ import time
1487
+
1488
+ execution_id_short = input.execution_id[:8]
1489
+
1490
+ activity.logger.info(
1491
+ "analytics_activity_started",
1492
+ extra={
1493
+ "execution_id": execution_id_short,
1494
+ "turn_number": input.turn_number,
1495
+ "tokens": input.result.get("usage", {}).get("total_tokens", 0),
1496
+ }
1497
+ )
1498
+
1499
+ try:
1500
+ # Initialize analytics service with required parameters
1501
+ control_plane_url = os.environ.get("CONTROL_PLANE_URL", "http://localhost:8000")
1502
+ api_key = os.environ.get("KUBIYA_API_KEY", "")
1503
+
1504
+ if not api_key:
1505
+ raise ValueError("KUBIYA_API_KEY environment variable not set")
1506
+
1507
+ analytics_service = AnalyticsService(
1508
+ control_plane_url=control_plane_url,
1509
+ api_key=api_key,
1510
+ )
1511
+
1512
+ # Convert dict result to RuntimeExecutionResult
1513
+ # The workflow passes a dict, but submit_runtime_analytics expects RuntimeExecutionResult
1514
+ result_obj = RuntimeExecutionResult(
1515
+ response=input.result.get("response", ""),
1516
+ usage=input.result.get("usage", {}),
1517
+ success=input.result.get("success", True),
1518
+ finish_reason=input.result.get("finish_reason", "stop"),
1519
+ tool_execution_messages=input.result.get("tool_messages", []),
1520
+ tool_messages=input.result.get("tool_messages", []),
1521
+ model=input.result.get("model", "unknown"),
1522
+ metadata=input.result.get("metadata", {}),
1523
+ error=input.result.get("error"),
1524
+ )
1525
+
1526
+ # Submit analytics
1527
+ try:
1528
+ await submit_runtime_analytics(
1529
+ result=result_obj,
1530
+ execution_id=input.execution_id,
1531
+ turn_number=input.turn_number,
1532
+ turn_start_time=input.turn_start_time,
1533
+ analytics_service=analytics_service,
1534
+ turn_end_time=time.time(),
1535
+ )
1536
+
1537
+ activity.logger.info(
1538
+ "analytics_activity_completed",
1539
+ extra={
1540
+ "execution_id": execution_id_short,
1541
+ "turn_number": input.turn_number,
1542
+ }
1543
+ )
1544
+
1545
+ return {
1546
+ "success": True,
1547
+ "execution_id": input.execution_id,
1548
+ "turn_number": input.turn_number,
1549
+ }
1550
+
1551
+ finally:
1552
+ # Cleanup HTTP client
1553
+ try:
1554
+ await analytics_service.aclose()
1555
+ except Exception as cleanup_error:
1556
+ activity.logger.warning(
1557
+ "analytics_service_cleanup_error",
1558
+ extra={
1559
+ "execution_id": execution_id_short,
1560
+ "error": str(cleanup_error),
1561
+ }
1562
+ )
1563
+
1564
+ except Exception as e:
1565
+ # Log but return success - analytics errors are non-critical
1566
+ error_msg = str(e)[:500]
1567
+ activity.logger.warning(
1568
+ "analytics_activity_failed",
1569
+ extra={
1570
+ "execution_id": execution_id_short,
1571
+ "turn_number": input.turn_number,
1572
+ "error": error_msg,
1573
+ "error_type": type(e).__name__,
1574
+ "note": "Analytics failed but execution continues normally"
1575
+ }
1576
+ )
1577
+
1578
+ return {
1579
+ "success": False,
1580
+ "execution_id": input.execution_id,
1581
+ "turn_number": input.turn_number,
1582
+ "error": error_msg,
1583
+ "error_type": type(e).__name__,
1584
+ }
1585
+