diff --git a/daiv/accounts/locale/pt/LC_MESSAGES/django.po b/daiv/accounts/locale/pt/LC_MESSAGES/django.po index 1fbaeda70..638447a5a 100644 --- a/daiv/accounts/locale/pt/LC_MESSAGES/django.po +++ b/daiv/accounts/locale/pt/LC_MESSAGES/django.po @@ -97,6 +97,9 @@ msgstr "" msgid "Your DAIV Sign-In Code" msgstr "O Seu Código de Autenticação DAIV" +msgid "New chat" +msgstr "Novo chat" + msgid "Dashboard" msgstr "Painel" diff --git a/daiv/automation/agent/base.py b/daiv/automation/agent/base.py index 629944181..c54289158 100644 --- a/daiv/automation/agent/base.py +++ b/daiv/automation/agent/base.py @@ -17,10 +17,12 @@ from core.constants import BOT_NAME from core.models import Provider, ProviderType from core.models import ThinkingLevelChoices as ThinkingLevel +from core.site_settings import site_settings logger = logging.getLogger("daiv.automation") if TYPE_CHECKING: + import httpx from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import BaseMessage from langgraph.checkpoint.base import BaseCheckpointSaver @@ -266,8 +268,44 @@ def _apply_insecure_http_clients(kw: dict, row: Provider.Cached) -> None: return import httpx - kw["http_client"] = httpx.Client(verify=False) # noqa: S501 # admin-opted-in via Provider.verify_ssl - kw["http_async_client"] = httpx.AsyncClient(verify=False) # noqa: S501 + # Carry over the resolved request timeout; a custom client otherwise falls back to + # httpx's 5s default, which would cut off real generations. + timeout = kw["timeout"] + kw["http_client"] = httpx.Client(verify=False, timeout=timeout) # noqa: S501 # admin-opted-in via Provider.verify_ssl + kw["http_async_client"] = httpx.AsyncClient(verify=False, timeout=timeout) # noqa: S501 + + +def _resolve_request_timeout(provider_type: ProviderType, timeout_seconds: float) -> httpx.Timeout | float: + """Return a per-request timeout in the shape each provider's langchain integration accepts. + + ``langchain-openai`` (OpenAI + OpenRouter) accepts an ``httpx.Timeout``, so we give the connect + phase its own short fuse while read/write/pool get the full budget. ``langchain-anthropic`` + (``timeout: float | None``) and ``langchain-google-genai`` (scalar seconds, converted to ms + internally) only accept a plain number, so those get a bare float. + """ + timeout_seconds = float(timeout_seconds) + if provider_type in _HTTPX_CLIENT_PROVIDER_TYPES: + import httpx + + return httpx.Timeout(timeout_seconds, connect=min(10.0, timeout_seconds)) + return timeout_seconds + + +def _apply_request_timeout_and_retries(kw: dict, provider_type: ProviderType) -> None: + """Bound how long a model call can hang and how many times it retries. + + Replaces the SDKs' own defaults (600s timeouts for OpenAI/Anthropic, unbounded for Google) + with the site-configured budget. Caller-supplied ``timeout`` / ``max_retries`` (e.g. the + tighter web_fetch budget) take precedence, but an explicit ``None`` is treated as unset so a + caller cannot silently reintroduce an unbounded default. + """ + if kw.get("max_retries") is None: + retries = site_settings.model_max_retries + # langchain-google-genai's max_retries counts total *attempts* (0 and 1 both mean "no + # retries"), so translate to keep "N retries" meaning uniform across providers. + kw["max_retries"] = retries + 1 if provider_type == ProviderType.GOOGLE_GENAI else retries + if kw.get("timeout") is None: + kw["timeout"] = _resolve_request_timeout(provider_type, site_settings.model_request_timeout_seconds) _BARE_NAME_HEURISTICS = ( @@ -421,6 +459,8 @@ def get_model_kwargs(*, resolved: ResolvedProvider, thinking_level: ThinkingLeve else: raise RuntimeError(f"Unknown provider_type {row.provider_type!r} on slug {row.slug!r}") + _apply_request_timeout_and_retries(kw, row.provider_type) + if not row.verify_ssl: _apply_insecure_http_clients(kw, row) diff --git a/daiv/automation/agent/graph.py b/daiv/automation/agent/graph.py index def4d6e36..ec3882617 100644 --- a/daiv/automation/agent/graph.py +++ b/daiv/automation/agent/graph.py @@ -26,6 +26,7 @@ ) from automation.agent.mcp.toolkits import MCPToolkit from automation.agent.middlewares.deferred_tools import deferred_tools_middleware, direct_mcp_tools +from automation.agent.middlewares.delegate_jobs import DelegateJobsMiddleware from automation.agent.middlewares.ensure_response import ensure_non_empty_response from automation.agent.middlewares.file_system import ( CUSTOM_TOOL_DESCRIPTIONS, @@ -310,6 +311,7 @@ async def create_daiv_agent( # source of write_todos and the harness profile excludes nothing here. TodoListMiddleware(system_prompt=dynamic_write_todos_system_prompt(bash_tool_enabled=_sandbox_enabled)), *([SlashCommandMiddleware(subagents=subagents)] if ctx.config.slash_commands.enabled else []), + *([DelegateJobsMiddleware()] if ctx.config.orchestration.enabled else []), *( [SandboxMiddleware(agent_root=agent_root, client=run_client, sandbox_backend=sandbox_backend)] if _sandbox_enabled diff --git a/daiv/automation/agent/middlewares/delegate_jobs.py b/daiv/automation/agent/middlewares/delegate_jobs.py new file mode 100644 index 000000000..ed3e501a0 --- /dev/null +++ b/daiv/automation/agent/middlewares/delegate_jobs.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import json +import logging +from typing import TYPE_CHECKING + +from django.urls import reverse + +from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse +from langchain_core.runnables import RunnableConfig # noqa: TCH002 — used in @tool signature at runtime +from langchain_core.tools import tool +from pydantic import BaseModel, Field +from sandbox_envs.services import aresolve_repo_envs +from sessions.models import MAX_SPAWN_DEPTH, Session, SessionOrigin +from sessions.services import MAX_DELEGATED_TARGETS, RepoTarget, asubmit_batch_runs + +from codebase.authorization import REPO_ACCESS_DENIED_MESSAGE, RepositoryAccessDenied, aassert_can_run +from codebase.conf import settings +from codebase.models import RepositoryCatalog + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + +logger = logging.getLogger("daiv.tools") + +DELEGATE_JOBS_NAME = "delegate_jobs" + + +class DelegateTarget(BaseModel): + repo_id: str = Field(description="Identifier of the target repository.") + ref: str | None = Field(default=None, description="Base branch/ref to start from; omit for the default branch.") + prompt: str = Field( + min_length=1, + description="Self-contained instruction for this repository. The leg runs in an isolated session " + "and sees only this text — not your conversation, the originating request, or anything else you " + "can see here — so include all the context it needs. End with the no-change convention: " + "'if this repository is unaffected, reply saying so and make no changes.'", + ) + + +DELEGATE_JOBS_DESCRIPTION = f"""\ +Delegate tailored sub-jobs to other repositories and return immediately. + +Each target runs as an independent single-repo job (own thread, own MR). Your turn ends after +delegating; you will be resumed automatically with a summary of every leg once they all finish. + +- `goal`: one-line description of the overall objective (used to title the batch). +- `targets`: 1-{MAX_DELEGATED_TARGETS} entries, each `{{repo_id, ref?, prompt}}` — `prompt` required per target. + +Returns JSON: +{{"batch_id", "delegated": [{{repo_id, ref, thread_id, session_url}}], "failed": [{{repo_id, error}}]}}.""" + + +DELEGATE_JOBS_SYSTEM_PROMPT = f"""\ +## Delegation tool `{DELEGATE_JOBS_NAME}` + +When a task spans other repositories, use `{DELEGATE_JOBS_NAME}` to fan tailored work out to them — +each target runs as an independent job. This is only for work in *other* repositories; for parallel +work inside this one, use the `task` tool (subagents) instead. If the task is contained to this +repository, ignore this tool. +After a call whose `delegated` list is non-empty, state your plan and end your turn — do NOT poll +or wait. You will be resumed with the consolidated results when every delegated leg finishes. +If `delegated` comes back empty, nothing was started and no resume will come — handle the failures +in this same turn instead of waiting. +Each leg runs in isolation and sees only the prompt you give it, so make every target's prompt +self-contained — include the context it needs and the no-change convention. +""" + + +def _error(message: str) -> str: + return json.dumps({"error": message}) + + +@tool(DELEGATE_JOBS_NAME, description=DELEGATE_JOBS_DESCRIPTION) +async def delegate_jobs_tool(goal: str, targets: list[DelegateTarget], config: RunnableConfig) -> str: + """Delegate per-repo sub-jobs; returns a JSON string (no state mutation).""" + thread_id = (config.get("configurable") or {}).get("thread_id") + if not thread_id: + return _error("delegate_jobs is only available inside a checkpointed run.") + + session = await Session.objects.select_related("user").filter(thread_id=thread_id).afirst() + if session is None: + return _error("Could not resolve the current session for delegation.") + if session.user_id is None: + return _error("Delegation requires an authenticated coordinator; this session has no user.") + + if not targets: + return _error("At least one target is required.") + if len(targets) > MAX_DELEGATED_TARGETS: + return _error(f"At most {MAX_DELEGATED_TARGETS} targets per delegate_jobs call.") + if session.spawn_depth >= MAX_SPAWN_DEPTH: + return _error(f"Delegation depth limit reached (MAX_SPAWN_DEPTH={MAX_SPAWN_DEPTH}).") + # An omitted ref and the default branch's explicit name are the same physical checkout, so + # resolve both through the synced catalog before comparing. + slugs = {session.repo_id} | {t.repo_id for t in targets} + default_branches: dict[str, str] = { + slug: branch + async for slug, branch in RepositoryCatalog.objects.filter( + provider=settings.CLIENT.value, slug__in=slugs + ).values_list("slug", "default_branch") + } + + def checkout(repo_id: str, ref: str | None) -> tuple[str, str]: + return (repo_id, ref or default_branches.get(repo_id, "")) + + coordinator_checkout = checkout(session.repo_id, session.ref) + if any(checkout(t.repo_id, t.ref) == coordinator_checkout for t in targets): + # In-repo parallelism belongs to subagents; a different ref on the same repo is a distinct + # checkout and may delegate. + return _error( + f"Cannot delegate to the coordinator's own checkout ({session.repo_id!r} on " + f"{session.ref or 'default branch'!r}). delegate_jobs fans work out to other " + "checkouts as independent jobs; for parallel work on this one, use the `task` " + "tool (subagents) instead." + ) + + seen: set[tuple[str, str]] = set() + for t in targets: + key = checkout(t.repo_id, t.ref) + if key in seen: + return _error(f"Duplicate target: {t.repo_id} on {t.ref or 'default branch'}.") + seen.add(key) + + user = session.user + + # Per-target authorization: aassert_can_run is all-or-nothing, so partition on the denied set. + denied: set[str] = set() + try: + await aassert_can_run(user, [t.repo_id for t in targets]) + except RepositoryAccessDenied as exc: + denied = set(exc.repo_ids) + + allowed = [t for t in targets if t.repo_id not in denied] + failed = [{"repo_id": rid, "error": REPO_ACCESS_DENIED_MESSAGE} for rid in sorted(denied)] + + batch_id: str | None = None + delegated: list[dict] = [] + if allowed: + repo_targets = [RepoTarget(repo_id=t.repo_id, ref=t.ref or "", prompt=t.prompt) for t in allowed] + # A raise here (OperationalError, revoked-access RepositoryAccessDenied, validation + # ValueError) must surface as the tool's JSON error contract, not a tool-node crash. + try: + repo_targets = await aresolve_repo_envs(user=user, repos=repo_targets, explicit_env_id=None) + result = await asubmit_batch_runs( + user=user, + prompt=goal, + repos=repo_targets, + trigger_type=SessionOrigin.DELEGATED_JOB, + parent_thread_id=thread_id, + spawn_depth=session.spawn_depth + 1, + ) + except Exception: # noqa: BLE001 + logger.exception("delegate_jobs: submission failed for thread=%s", thread_id) + return _error("Delegation submission failed; no sub-jobs were started.") + batch_id = str(result.batch_id) + delegated = [ + { + "repo_id": run.repo_id, + "ref": run.ref, + "thread_id": str(run.session_id), + "session_url": reverse("session_detail", kwargs={"thread_id": run.session_id}), + } + for run in result.runs + ] + failed.extend({"repo_id": f.repo_id, "error": f.error} for f in result.failed) + + payload: dict = {"batch_id": batch_id, "delegated": delegated, "failed": failed} + if not delegated: + payload["note"] = ( + "No sub-jobs are running. Do not end your turn to wait for a resume — none will come; " + "handle the failures now." + ) + return json.dumps(payload, ensure_ascii=False) + + +class DelegateJobsMiddleware(AgentMiddleware): + """Bind the delegate_jobs tool and inject its usage note. Added to the agent when + ``orchestration.enabled`` is set — on by default; a repo opts out with + ``orchestration.enabled: false``. + """ + + def __init__(self) -> None: + self.tools = [delegate_jobs_tool] + + async def awrap_model_call( + self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]] + ) -> ModelResponse: + system_prompt = (request.system_prompt + "\n\n") if request.system_prompt else "" + system_prompt += DELEGATE_JOBS_SYSTEM_PROMPT + return await handler(request.override(system_prompt=system_prompt)) diff --git a/daiv/automation/agent/middlewares/web_fetch.py b/daiv/automation/agent/middlewares/web_fetch.py index f04d91bad..381bb6714 100644 --- a/daiv/automation/agent/middlewares/web_fetch.py +++ b/daiv/automation/agent/middlewares/web_fetch.py @@ -27,6 +27,11 @@ WEB_FETCH_NAME = "web_fetch" +# The summariser runs a small/fast model over already-bounded page content +# (``web_fetch_max_content_chars``), and the tool degrades to returning the raw +# content on failure — so keep its budget tighter than the global model timeout. +WEB_FETCH_MODEL_TIMEOUT_SECONDS = 60 + WEB_FETCH_TOOL_DESCRIPTION = """\ Fetch content from a specified URL and process it using an AI model. @@ -232,17 +237,30 @@ async def web_fetch_tool( if not prompt.strip() or site_settings.web_fetch_model_name is None: return f"Contents of {url}:\n{content}" + # Building the model is pure construction, so a failure here is a permanent provider + # misconfiguration that breaks every call — log at ERROR (Sentry) rather than degrade quietly. + try: + model = BaseAgent.get_model(model=site_settings.web_fetch_model_name, timeout=WEB_FETCH_MODEL_TIMEOUT_SECONDS) + except Exception as e: + logger.exception( + "web_fetch summariser model %r could not be built; check the provider configuration.", + site_settings.web_fetch_model_name, + ) + return f"web_fetch summariser unavailable ({e}); returning raw content instead.\nContents of {url}:\n{content}" + + messages = [ + SystemMessage( + content=( + "You process web pages for users. Use the page content to answer the user's prompt.\n" + "Be concise. If the content doesn't contain the answer, say so." + ) + ), + HumanMessage(content=f"URL: {url}\n\n\n{content}\n\n\nPrompt:\n{prompt}"), + ] + + # Invocation failures are transient (timeout, network, provider error) — degrade gracefully + # to the raw page content so a blip doesn't abort the agent's turn. try: - model = BaseAgent.get_model(model=site_settings.web_fetch_model_name) - messages = [ - SystemMessage( - content=( - "You process web pages for users. Use the page content to answer the user's prompt.\n" - "Be concise. If the content doesn't contain the answer, say so." - ) - ), - HumanMessage(content=f"URL: {url}\n\n\n{content}\n\n\nPrompt:\n{prompt}"), - ] response = await model.ainvoke(messages) response_text = str(getattr(response, "content", response)) _set_cached_response(url=url, prompt=prompt, response=response_text) diff --git a/daiv/automation/agent/skills/orchestrate/SKILL.md b/daiv/automation/agent/skills/orchestrate/SKILL.md new file mode 100644 index 000000000..dd92e7dbf --- /dev/null +++ b/daiv/automation/agent/skills/orchestrate/SKILL.md @@ -0,0 +1,43 @@ +--- +name: orchestrate +description: Use when you are a coordinator on a coordination repo turning one request — a ticket, issue, security advisory, or direct instruction — into tailored per-repo sub-jobs via delegate_jobs, then reporting the consolidated outcome once every leg finishes. +--- + +# Orchestrate cross-repo work + +You run on a **coordination repository**. Your job is to turn one request — a ticket, issue, +security advisory, or direct instruction — into tailored work across other repositories, then +report the combined outcome. + +> **Precondition.** This skill depends on the `delegate_jobs` tool, which is bound by default but can +> be disabled per-repository via `orchestration.enabled: false` in `.daiv.yml`. If `delegate_jobs` is +> not in your available tools, orchestration has been disabled here — say so plainly and stop; do not +> try to emulate delegation by other means. + +## Workflow + +1. **Triage.** Read the request using your attached MCP tools. Consult `AGENTS.md` in this repo for + the repository directory and routing rules. Decide which repositories are affected and what each + one must do. +2. **Delegate.** Call `delegate_jobs(goal, targets)` once, with a tailored `prompt` per target. + - Each leg runs in an isolated session and sees only the prompt you give it — not this + conversation, the originating request, or your tool outputs. Put everything a leg needs + directly in its prompt. + - Include the no-change convention in each prompt: *"if this repository is unaffected, reply + saying so and make no changes."* +3. **End your turn.** State your delegation plan and stop. Do **not** poll or wait — you will be + resumed automatically once every leg finishes. +4. **Report.** On resume you receive a summary of all legs (status, MR links, replies). Verify the + outcome, compose the consolidated result, and report it back to wherever the request originated + (e.g. comment on the ticket or issue) using your MCP tools. +5. **Follow up (optional).** If a sequenced change is needed (e.g. adapt repo B against repo A's + MR), delegate another batch — you will be resumed again. + +## Limits + +- Up to 10 targets per `delegate_jobs` call. +- Delegation depth is capped; a leg cannot itself delegate beyond the configured chain depth. +- You can only delegate as an authenticated coordinator; targets you lack write access to are + reported back to you as failures rather than run. + +Routing rules specific to your setup belong in this repo's `.agents/AGENTS.md`, not this skill. diff --git a/daiv/chat/api/streaming.py b/daiv/chat/api/streaming.py index ec61add22..28cf3eea6 100644 --- a/daiv/chat/api/streaming.py +++ b/daiv/chat/api/streaming.py @@ -10,6 +10,7 @@ from django.utils import timezone from ag_ui.core.events import BaseEvent, CustomEvent, EventType, RunErrorEvent +from asgiref.sync import sync_to_async from copilotkit import LangGraphAGUIAgent from langgraph.store.memory import InMemoryStore from sessions.locks import SessionLock @@ -403,3 +404,11 @@ async def events(self) -> AsyncIterator[BaseEvent]: await SessionLock.release(self.thread_id, self.run_id) except Exception: logger.exception("chat: failed to release run slot for thread_id=%s", self.thread_id) + # Chat runs emit no run_finished, so a delegated-batch continuation that completed + # during this turn parks QUEUED; release it now that the slot is free. + try: + from sessions.signals import release_next_queued + + await sync_to_async(release_next_queued, thread_sensitive=True)(self.thread_id) + except Exception: + logger.exception("chat: failed to release queued runs for thread_id=%s", self.thread_id) diff --git a/daiv/codebase/repo_config.py b/daiv/codebase/repo_config.py index 696a63f38..56acc1a21 100644 --- a/daiv/codebase/repo_config.py +++ b/daiv/codebase/repo_config.py @@ -83,6 +83,14 @@ class Memory(BaseModel): ) +class Orchestration(BaseModel): + """ + Orchestration configuration. + """ + + enabled: bool = Field(default=True, description="Bind the delegate_jobs tool for agent runs on this repository.") + + class AgentModelConfig(BaseModel): """ Model configuration for the DAIV agent. @@ -191,6 +199,9 @@ class RepositoryConfig(BaseModel): default_factory=IssueAddressing, description="Configure issue addressing features." ) memory: Memory = Field(default_factory=Memory, description="Configure learned repository memory features.") + orchestration: Orchestration = Field( + default_factory=Orchestration, description="Configure orchestrated delegation features." + ) models: Models = Field(default_factory=Models, description="Configure model settings for agents.") def is_user_allowed(self, username: str) -> bool: diff --git a/daiv/core/forms.py b/daiv/core/forms.py index 7888a4960..d81bc6e6d 100644 --- a/daiv/core/forms.py +++ b/daiv/core/forms.py @@ -204,6 +204,8 @@ class Meta: "agent_explore_model_name", "agent_explore_fallback_model_name", "agent_recursion_limit", + "model_request_timeout_seconds", + "model_max_retries", "diff_to_metadata_model_name", "diff_to_metadata_fallback_model_name", "titling_model_name", diff --git a/daiv/core/migrations/0015_siteconfiguration_model_max_retries_and_more.py b/daiv/core/migrations/0015_siteconfiguration_model_max_retries_and_more.py new file mode 100644 index 000000000..b57df5886 --- /dev/null +++ b/daiv/core/migrations/0015_siteconfiguration_model_max_retries_and_more.py @@ -0,0 +1,32 @@ +# Generated by Django 6.0.7 on 2026-07-13 20:52 + +import django.core.validators +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [("core", "0014_siteconfiguration_memory_flush_age_and_budget_bounds")] + + operations = [ + migrations.AddField( + model_name="siteconfiguration", + name="model_max_retries", + field=models.PositiveIntegerField( + blank=True, + help_text="Retries for a failed LLM API call before the error propagates. Applies to every provider.", + null=True, + verbose_name="model max retries", + ), + ), + migrations.AddField( + model_name="siteconfiguration", + name="model_request_timeout_seconds", + field=models.PositiveIntegerField( + blank=True, + help_text="Per-request timeout for LLM API calls, in seconds. Applies to every provider.", + null=True, + validators=[django.core.validators.MinValueValidator(1)], + verbose_name="model request timeout", + ), + ), + ] diff --git a/daiv/core/models.py b/daiv/core/models.py index 8d620b67f..45a0e2ea3 100644 --- a/daiv/core/models.py +++ b/daiv/core/models.py @@ -246,6 +246,19 @@ class SiteConfiguration(models.Model): agent_recursion_limit = models.PositiveIntegerField( _("recursion limit"), blank=True, null=True, help_text=_("Maximum recursion depth for agent loops.") ) + model_request_timeout_seconds = models.PositiveIntegerField( + _("model request timeout"), + blank=True, + null=True, + validators=[MinValueValidator(1)], + help_text=_("Per-request timeout for LLM API calls, in seconds. Applies to every provider."), + ) + model_max_retries = models.PositiveIntegerField( + _("model max retries"), + blank=True, + null=True, + help_text=_("Retries for a failed LLM API call before the error propagates. Applies to every provider."), + ) # -- Commit & PR Writer -- diff_to_metadata_model_name = models.CharField( @@ -531,7 +544,7 @@ class SiteConfiguration(models.Model): FieldGroup( key="agent", title=_("Agent"), - match=("agent_*", "suggest_context_file_enabled"), + match=("agent_*", "suggest_context_file_enabled", "model_request_timeout_seconds", "model_max_retries"), icon="agent", category="AI tasks", ), diff --git a/daiv/core/site_settings.py b/daiv/core/site_settings.py index d6ef3d22d..65798e7ac 100644 --- a/daiv/core/site_settings.py +++ b/daiv/core/site_settings.py @@ -68,6 +68,10 @@ def _build_field_defaults() -> dict[str, Any]: "agent_explore_model_name": ModelName.CLAUDE_HAIKU_4_5, "agent_explore_fallback_model_name": ModelName.GPT_5_4_MINI, "agent_recursion_limit": 500, + # Matches the OpenAI/Anthropic SDK default; anything shorter cuts off long non-streaming + # generations (the agent invokes models without streaming). + "model_request_timeout_seconds": 600, + "model_max_retries": 1, "suggest_context_file_enabled": True, # Diff to Metadata "diff_to_metadata_model_name": ModelName.GPT_5_4_MINI, diff --git a/daiv/jobs/api/schemas.py b/daiv/jobs/api/schemas.py index 7a5cef1ce..6a6c4df4a 100644 --- a/daiv/jobs/api/schemas.py +++ b/daiv/jobs/api/schemas.py @@ -12,6 +12,9 @@ class RepoSubmitItem(Schema): repo_id: str = Field(min_length=1) ref: str | None = None + prompt: str | None = Field( + default=None, description="Optional per-repo instruction; overrides the batch prompt for this repo." + ) class JobSubmitRequest(Schema): diff --git a/daiv/jobs/api/views.py b/daiv/jobs/api/views.py index b0fa98427..811c9f4c3 100644 --- a/daiv/jobs/api/views.py +++ b/daiv/jobs/api/views.py @@ -75,7 +75,7 @@ async def submit_job(request: HttpRequest, payload: JobSubmitRequest): return 400, {"detail": str(err)} explicit_env_id = str(env.id) if env else None - targets = [RepoTarget(repo_id=spec.repo_id, ref=spec.ref or "") for spec in payload.repos] + targets = [RepoTarget(repo_id=spec.repo_id, ref=spec.ref or "", prompt=spec.prompt) for spec in payload.repos] targets = await aresolve_repo_envs(user=request.auth, repos=targets, explicit_env_id=explicit_env_id) result = await asubmit_batch_runs( user=request.auth, diff --git a/daiv/mcp_server/server.py b/daiv/mcp_server/server.py index 390c7f471..c56630301 100644 --- a/daiv/mcp_server/server.py +++ b/daiv/mcp_server/server.py @@ -104,6 +104,9 @@ class RepoSubmitSpec(BaseModel): ref: str | None = Field( default=None, description="Git reference (branch name or commit SHA). None / empty string = default branch." ) + prompt: str | None = Field( + default=None, description="Optional per-repo instruction; overrides the batch prompt for this repo." + ) @mcp.tool() @@ -119,7 +122,8 @@ async def submit_job( " to commit, push, create a branch, or open a merge/pull request; those steps run" " automatically after the job, and DAIV generates the branch name, commit message," " and MR/PR title/description itself. Be specific: include file paths, function" - " names, or error messages. The same prompt runs independently against each repository." + " names, or error messages. The same prompt runs independently against each repository" + " unless a per-repo override is set via repos[].prompt." ) ), ], @@ -252,7 +256,7 @@ async def submit_job( if env_row is not None: explicit_env_id = str(env_row.id) - targets = [RepoTarget(repo_id=s.repo_id, ref=s.ref or "") for s in specs] + targets = [RepoTarget(repo_id=s.repo_id, ref=s.ref or "", prompt=s.prompt) for s in specs] targets = await aresolve_repo_envs(user=mcp_user, repos=targets, explicit_env_id=explicit_env_id) result = await asubmit_batch_runs( user=mcp_user, diff --git a/daiv/notifications/signals.py b/daiv/notifications/signals.py index b4113fa94..63730f6b9 100644 --- a/daiv/notifications/signals.py +++ b/daiv/notifications/signals.py @@ -279,6 +279,11 @@ def on_run_finished(sender, run, **kwargs) -> None: if run.trigger_type in (SessionOrigin.webhooks() | {SessionOrigin.CHAT}): return + # Coordinator continuation runs are internal plumbing; the legs' batch/per-run + # notification below is the user-facing signal. + if run.continuation_of_batch_id is not None: + return + if run.batch_id is not None: siblings = Run.objects.by_batch(run.batch_id) total = siblings.count() diff --git a/daiv/sessions/management/commands/release_orphan_queued_sessions.py b/daiv/sessions/management/commands/release_orphan_queued_sessions.py index d3c3c2c24..8987a8641 100644 --- a/daiv/sessions/management/commands/release_orphan_queued_sessions.py +++ b/daiv/sessions/management/commands/release_orphan_queued_sessions.py @@ -4,10 +4,10 @@ from django.core.management.base import BaseCommand from django.db import IntegrityError -from django.db.models import Q +from django.db.models import Exists, OuterRef from sessions.models import Run, RunStatus -from sessions.signals import _enqueue_queued_run +from sessions.signals import DISPATCH_FAILED_PREFIX, _enqueue_queued_run logger = logging.getLogger("daiv.sessions") @@ -16,18 +16,32 @@ class Command(BaseCommand): help = ( "Release QUEUED Runs whose session has no active (READY/RUNNING) sibling. " "Mitigates a rare TOCTOU loss where the dispatcher missed a terminal transition " - "or the row was created QUEUED but never picked up." + "or the row was created QUEUED but never picked up. Delegated-batch continuation " + "runs that FAILED before ever starting (dispatch failure, no linked task) are " + "re-queued first so the same pass can release them." ) def handle(self, *args, **options): - active_sessions = set( - Run.objects.filter(status__in=[RunStatus.READY, RunStatus.RUNNING]).values_list("session_id", flat=True) - ) + # A continuation is a batch's one shot at resuming its coordinator + # (run_one_continuation_per_batch), so a never-started FAILED row must be retried. + # A row that reached the broker (task_result_id set, or the link_failed prefix) is not. + requeued = Run.objects.filter( + continuation_of_batch_id__isnull=False, + status=RunStatus.FAILED, + task_result_id__isnull=True, + error_message__startswith=DISPATCH_FAILED_PREFIX, + ).update(status=RunStatus.QUEUED, error_message="", finished_at=None, started_at=None) orphans = ( Run.objects .filter(status=RunStatus.QUEUED) - .filter(~Q(session_id__in=active_sessions)) + .exclude( + Exists( + Run.objects.filter( + session_id=OuterRef("session_id"), status__in=[RunStatus.READY, RunStatus.RUNNING] + ) + ) + ) .order_by("session_id", "created_at") ) @@ -41,8 +55,8 @@ def handle(self, *args, **options): try: claimed = Run.objects.filter(pk=run.pk, status=RunStatus.QUEUED).update(status=RunStatus.READY) except IntegrityError: - # A concurrent submission claimed the session between our snapshot of - # active_sessions and this CAS; leave the row QUEUED for a future pass. + # A concurrent submission claimed the session between the orphan scan and + # this CAS; leave the row QUEUED for a future pass. skipped += 1 continue if claimed != 1: @@ -60,7 +74,7 @@ def handle(self, *args, **options): else: errored += 1 - summary = f"Released: {released}, skipped: {skipped}, errored: {errored}" + summary = f"Released: {released}, requeued continuations: {requeued}, skipped: {skipped}, errored: {errored}" if errored: self.stdout.write(self.style.WARNING(f"{summary} — see logs; broker may be unavailable.")) else: diff --git a/daiv/sessions/migrations/0005_delegate_jobs_fields.py b/daiv/sessions/migrations/0005_delegate_jobs_fields.py new file mode 100644 index 000000000..a2ef7d26a --- /dev/null +++ b/daiv/sessions/migrations/0005_delegate_jobs_fields.py @@ -0,0 +1,143 @@ +# Generated by Django 6.0.7 on 2026-07-09 10:25 + +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_sessions", "0004_runenvelope"), + ("django_tasks_database", "0019_rename_django_task_new_ordering_idx_tasks_db_new_ordering_idx_and_more"), + ("sandbox_envs", "0006_drop_network_enabled"), + ("schedules", "0016_alter_scheduledjob_agent_thinking_level_and_more"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.RemoveConstraint(model_name="run", name="run_one_active_per_session"), + migrations.RemoveConstraint(model_name="run", name="run_trigger_type_valid"), + migrations.RemoveConstraint(model_name="session", name="session_origin_valid"), + migrations.AddField( + model_name="run", + name="continuation_of_batch_id", + field=models.UUIDField( + blank=True, + help_text="Set on the coordinator continuation run enqueued when this batch turned terminal.", + null=True, + verbose_name="continuation of batch", + ), + ), + migrations.AddField( + model_name="session", + name="parent_thread_id", + field=models.CharField( + blank=True, + default=None, + help_text="Coordinator thread that delegated this session's first run via delegate_jobs; empty for top-level sessions.", + max_length=64, + null=True, + verbose_name="parent thread ID", + ), + ), + migrations.AddField( + model_name="session", + name="spawn_depth", + field=models.PositiveSmallIntegerField( + default=0, + help_text="0 = human/system-triggered; +1 per delegate_jobs hop. Capped by MAX_SPAWN_DEPTH.", + verbose_name="spawn depth", + ), + ), + migrations.AlterField( + model_name="run", + name="trigger_type", + field=models.CharField( + choices=[ + ("chat", "Chat"), + ("api_job", "API Run"), + ("mcp_job", "MCP Run"), + ("schedule", "Scheduled Run"), + ("ui_job", "UI Run"), + ("issue_webhook", "Issue Webhook"), + ("mr_webhook", "MR/PR Webhook"), + ("delegated_job", "Delegated Run"), + ], + max_length=20, + verbose_name="trigger type", + ), + ), + migrations.AlterField( + model_name="session", + name="origin", + field=models.CharField( + choices=[ + ("chat", "Chat"), + ("api_job", "API Run"), + ("mcp_job", "MCP Run"), + ("schedule", "Scheduled Run"), + ("ui_job", "UI Run"), + ("issue_webhook", "Issue Webhook"), + ("mr_webhook", "MR/PR Webhook"), + ("delegated_job", "Delegated Run"), + ], + max_length=20, + verbose_name="origin", + ), + ), + migrations.AddConstraint( + model_name="run", + constraint=models.UniqueConstraint( + condition=models.Q( + ("status__in", ["READY", "RUNNING"]), ("trigger_type__in", ["api_job", "mcp_job", "delegated_job"]) + ), + fields=("session",), + name="run_one_active_per_session", + ), + ), + migrations.AddConstraint( + model_name="run", + constraint=models.UniqueConstraint( + condition=models.Q(("continuation_of_batch_id__isnull", False)), + fields=("continuation_of_batch_id",), + name="run_one_continuation_per_batch", + ), + ), + migrations.AddConstraint( + model_name="run", + constraint=models.CheckConstraint( + condition=models.Q(( + "trigger_type__in", + [ + "chat", + "api_job", + "mcp_job", + "schedule", + "ui_job", + "issue_webhook", + "mr_webhook", + "delegated_job", + ], + )), + name="run_trigger_type_valid", + ), + ), + migrations.AddConstraint( + model_name="session", + constraint=models.CheckConstraint( + condition=models.Q(( + "origin__in", + [ + "chat", + "api_job", + "mcp_job", + "schedule", + "ui_job", + "issue_webhook", + "mr_webhook", + "delegated_job", + ], + )), + name="session_origin_valid", + ), + ), + ] diff --git a/daiv/sessions/migrations/0006_session_spawn_depth_cap.py b/daiv/sessions/migrations/0006_session_spawn_depth_cap.py new file mode 100644 index 000000000..14563445e --- /dev/null +++ b/daiv/sessions/migrations/0006_session_spawn_depth_cap.py @@ -0,0 +1,34 @@ +# Generated by Django 6.0.7 on 2026-07-09 13:41 + +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("agent_sessions", "0005_delegate_jobs_fields"), + ("sandbox_envs", "0006_drop_network_enabled"), + ("schedules", "0016_alter_scheduledjob_agent_thinking_level_and_more"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ] + + operations = [ + migrations.AlterField( + model_name="session", + name="parent_thread_id", + field=models.CharField( + blank=True, + default=None, + help_text="Coordinator thread that delegated this session's first run via delegate_jobs; unset (NULL) for top-level sessions.", + max_length=64, + null=True, + verbose_name="parent thread ID", + ), + ), + migrations.AddConstraint( + model_name="session", + constraint=models.CheckConstraint( + condition=models.Q(("spawn_depth__lte", 2)), name="session_spawn_depth_within_cap" + ), + ), + ] diff --git a/daiv/sessions/models.py b/daiv/sessions/models.py index c894f402a..270f87613 100644 --- a/daiv/sessions/models.py +++ b/daiv/sessions/models.py @@ -50,6 +50,7 @@ class SessionOrigin(models.TextChoices): UI_JOB = "ui_job", _("UI Run") ISSUE_WEBHOOK = "issue_webhook", _("Issue Webhook") MR_WEBHOOK = "mr_webhook", _("MR/PR Webhook") + DELEGATED_JOB = "delegated_job", _("Delegated Run") @classmethod def webhooks(cls) -> frozenset[str]: @@ -58,8 +59,15 @@ def webhooks(cls) -> frozenset[str]: @classmethod def prompt_driven(cls) -> frozenset[str]: - """Job triggers created from an explicit user prompt (API / MCP / UI batch submits).""" - return frozenset({cls.API_JOB, cls.MCP_JOB, cls.UI_JOB}) + """Job triggers created from an explicit user prompt (API / MCP / UI / delegated batch submits).""" + return frozenset({cls.API_JOB, cls.MCP_JOB, cls.UI_JOB, cls.DELEGATED_JOB}) + + +#: Maximum delegation depth (``Session.spawn_depth``). 0 = human/system-triggered; each +#: ``delegate_jobs`` hop adds 1. Enforced both at the tool boundary (``delegate_jobs``) and by +#: the ``session_spawn_depth_within_cap`` DB CheckConstraint below so no other ``asubmit_batch_runs`` +#: caller can stamp a leg past the fuse. +MAX_SPAWN_DEPTH = 2 class EnvelopeStatus(models.TextChoices): @@ -139,6 +147,22 @@ class Session(models.Model): ) issue_iid = models.PositiveIntegerField(_("issue IID"), null=True, blank=True) merge_request_iid = models.PositiveIntegerField(_("merge request IID"), null=True, blank=True) + parent_thread_id = models.CharField( # noqa: DJ001 — mirrors thread_id's shape; NULL = top-level + _("parent thread ID"), + max_length=64, + null=True, + blank=True, + default=None, + help_text=_( + "Coordinator thread that delegated this session's first run via delegate_jobs; " + "unset (NULL) for top-level sessions." + ), + ) + spawn_depth = models.PositiveSmallIntegerField( + _("spawn depth"), + default=0, + help_text=_("0 = human/system-triggered; +1 per delegate_jobs hop. Capped by MAX_SPAWN_DEPTH."), + ) # Unified execution lock. NULL means "free slot"; any non-NULL value is the # holder id (AG-UI run_id for chat turns, str(Run.pk) for background runs). @@ -176,6 +200,12 @@ class Meta: | models.Q(agent_thinking_level__in=ThinkingLevelChoices.values), name="session_agent_thinking_level_valid", ), + # Delegation-depth fuse: the ``delegate_jobs`` tool caps at ``MAX_SPAWN_DEPTH``, but + # ``asubmit_batch_runs(spawn_depth=...)`` is an open int — pin the ceiling at the DB so + # no future caller can stamp a leg past the recursion fuse. + models.CheckConstraint( + condition=models.Q(spawn_depth__lte=MAX_SPAWN_DEPTH), name="session_spawn_depth_within_cap" + ), ] def __str__(self) -> str: @@ -271,6 +301,12 @@ class Run(models.Model): total_tokens = models.PositiveIntegerField(_("total tokens"), null=True, blank=True) cost_usd = models.DecimalField(_("cost (USD)"), max_digits=10, decimal_places=6, null=True, blank=True) usage_by_model = models.JSONField(_("usage by model"), null=True, blank=True) + continuation_of_batch_id = models.UUIDField( + _("continuation of batch"), + null=True, + blank=True, + help_text=_("Set on the coordinator continuation run enqueued when this batch turned terminal."), + ) created_at = models.DateTimeField(_("created at"), default=timezone.now, editable=False) started_at = models.DateTimeField(_("started at"), null=True, blank=True) @@ -289,19 +325,28 @@ class Meta: models.Index(fields=["user", "-created_at"], name="run_user_created_idx"), ] constraints = [ - # At most one active (READY or RUNNING) API/MCP run per session. QUEUED is + # At most one active (READY or RUNNING) API/MCP/delegated-job run per session. QUEUED is # intentionally outside the constraint so FIFO siblings can stack; webhook # triggers share deterministic sessions and are intentionally excluded. # Exact port of activity_one_active_per_thread. The status/trigger_type # literals here must equal {RunStatus.READY, RunStatus.RUNNING} and - # {SessionOrigin.API_JOB, SessionOrigin.MCP_JOB} — Django serializes + # {SessionOrigin.API_JOB, SessionOrigin.MCP_JOB, SessionOrigin.DELEGATED_JOB} — Django serializes # constraints with literals, so a drift is caught by # ``test_active_constraint_literals_match_enums``. models.UniqueConstraint( fields=["session"], - condition=models.Q(status__in=["READY", "RUNNING"], trigger_type__in=["api_job", "mcp_job"]), + condition=models.Q( + status__in=["READY", "RUNNING"], trigger_type__in=["api_job", "mcp_job", "delegated_job"] + ), name="run_one_active_per_session", ), + # At most one coordinator-continuation run per delegated batch — the + # winner-election mechanism for the resume signal. + models.UniqueConstraint( + fields=["continuation_of_batch_id"], + condition=models.Q(continuation_of_batch_id__isnull=False), + name="run_one_continuation_per_batch", + ), # DB-level enum enforcement (``choices=`` alone is not enforced, and # ``.aupdate()``/raw writes bypass field validation). models.CheckConstraint( diff --git a/daiv/sessions/services.py b/daiv/sessions/services.py index c4bac646c..f054a6845 100644 --- a/daiv/sessions/services.py +++ b/daiv/sessions/services.py @@ -28,12 +28,15 @@ logger = logging.getLogger("daiv.sessions") +MAX_DELEGATED_TARGETS = 10 + @dataclass(frozen=True) class RepoTarget: repo_id: str ref: str = "" sandbox_environment_id: str | None = None + prompt: str | None = None # per-target override; falls back to the batch prompt @dataclass(frozen=True) @@ -72,6 +75,8 @@ async def aget_or_create_session( scheduled_job=None, issue_iid: int | None = None, merge_request_iid: int | None = None, + parent_thread_id: str | None = None, + spawn_depth: int = 0, ) -> Session: """Idempotent session bootstrap keyed on thread_id. First caller sets origin and context; later callers just bump last_active_at (a webhook session later @@ -92,6 +97,8 @@ async def aget_or_create_session( "scheduled_job": scheduled_job, "issue_iid": issue_iid, "merge_request_iid": merge_request_iid, + "parent_thread_id": parent_thread_id, + "spawn_depth": spawn_depth, }, ) if not created: @@ -121,6 +128,9 @@ async def acreate_run( title: str = "", sandbox_environment_id: str | None = None, status: str = RunStatus.READY, + parent_thread_id: str | None = None, + spawn_depth: int = 0, + continuation_of_batch_id: uuid.UUID | None = None, ) -> Run: """Async: create a Session (idempotent) then a Run linked to it. @@ -147,6 +157,8 @@ async def acreate_run( scheduled_job=scheduled_job, issue_iid=issue_iid, merge_request_iid=merge_request_iid, + parent_thread_id=parent_thread_id, + spawn_depth=spawn_depth, ) return await Run.objects.acreate( session=session, @@ -166,6 +178,7 @@ async def acreate_run( sandbox_environment_id=sandbox_environment_id, merge_request_iid=merge_request_iid, mention_comment_id=mention_comment_id, + continuation_of_batch_id=continuation_of_batch_id, ) @@ -209,6 +222,8 @@ async def asubmit_batch_runs( scheduled_job: ScheduledJob | None = None, external_username: str = "", thread_id: str | None = None, + parent_thread_id: str | None = None, + spawn_depth: int = 0, ) -> BatchSubmitResult: """Enqueue N ``run_job_task`` instances sharing a ``batch_id``; record N ``Run`` rows. @@ -217,6 +232,10 @@ async def asubmit_batch_runs( Best-effort: any per-repo exception (enqueue failure or post-enqueue run-creation failure) lands in ``result.failed`` while siblings continue. + + All Run rows are created before any task is enqueued, so ``run_finished`` receivers + that scan the batch (coordinator resume, rollup notification) never observe a + partially-created sibling set. """ _validate(repos) if user is not None: @@ -238,7 +257,10 @@ async def asubmit_batch_runs( if trigger_type == SessionOrigin.SCHEDULE and scheduled_job is not None: schedule_run_base = await Run.objects.filter(session__scheduled_job=scheduled_job).acount() - async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: + def _effective_prompt(target: RepoTarget) -> str: + return target.prompt or prompt + + async def _create_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: effective_thread_id = thread_id or str(uuid.uuid4()) run_title = "" @@ -249,7 +271,7 @@ async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: "trigger_type": trigger_type, "repo_id": target.repo_id, "ref": target.ref, - "prompt": prompt, + "prompt": _effective_prompt(target), "agent_model": agent_model, "agent_thinking_level": agent_thinking_level, "scheduled_job": scheduled_job, @@ -260,6 +282,8 @@ async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: "thread_id": effective_thread_id, "title": run_title, "sandbox_environment_id": target.sandbox_environment_id, + "parent_thread_id": parent_thread_id, + "spawn_depth": spawn_depth, } # Claim the session atomically by trying to create a READY row. The partial @@ -267,7 +291,7 @@ async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: # when a sibling (READY/RUNNING) is already active on this session — in that # case we fall back to QUEUED, no task enqueue. try: - run = await acreate_run(**common_kwargs, task_result_id=None, status=RunStatus.READY) + return await acreate_run(**common_kwargs, task_result_id=None, status=RunStatus.READY) except IntegrityError: try: return await acreate_run(**common_kwargs, task_result_id=None, status=RunStatus.QUEUED) @@ -284,14 +308,15 @@ async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: repo_id=target.repo_id, ref=target.ref, error=f"RunCreationFailed: {type(err).__name__}: {err}" ) + async def _enqueue_one(run: Run, target: RepoTarget) -> Run | BatchSubmitFailure: try: task = await run_job_task.aenqueue( repo_id=target.repo_id, - prompt=prompt, + prompt=_effective_prompt(target), ref=target.ref or None, agent_model=agent_model or None, agent_thinking_level=agent_thinking_level or None, - thread_id=effective_thread_id, + thread_id=str(run.session_id), sandbox_environment_id=target.sandbox_environment_id, run_id=str(run.pk), user_id=user.id if user is not None else None, @@ -312,9 +337,6 @@ async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: logger.exception( "submit_batch_runs: failed to link task_result_id=%s to run=%s (orphan task will run)", task.id, run.pk ) - # Surface in error_message that the agent may run to completion (push a - # commit / open an MR) while this row shows FAILED — the work is real but - # uncapturable because nothing links back to it. await _mark_failed_and_advance( run, prefix=LINK_FAILED_PREFIX, err=save_err, previous_status=RunStatus.READY ) @@ -324,8 +346,20 @@ async def _submit_one(idx: int, target: RepoTarget) -> Run | BatchSubmitFailure: return run # return_exceptions=True guards against BaseException (CancelledError, etc.) aborting the - # whole batch; _submit_one already catches Exception itself. - outcomes = await asyncio.gather(*[_submit_one(i, t) for i, t in enumerate(repos)], return_exceptions=True) + # whole batch; the helpers already catch Exception themselves. + created = await asyncio.gather(*[_create_one(i, t) for i, t in enumerate(repos)], return_exceptions=True) + + # QUEUED fallbacks are not enqueued here — the dispatcher releases them later. + to_enqueue = [ + (i, outcome, target) + for i, (target, outcome) in enumerate(zip(repos, created, strict=True)) + if not isinstance(outcome, BatchSubmitFailure | BaseException) and outcome.status == RunStatus.READY + ] + enqueued = await asyncio.gather(*[_enqueue_one(r, t) for _, r, t in to_enqueue], return_exceptions=True) + + outcomes: list = list(created) + for (i, _, _), outcome in zip(to_enqueue, enqueued, strict=True): + outcomes[i] = outcome runs: list[Run] = [] failed: list[BatchSubmitFailure] = [] diff --git a/daiv/sessions/signals.py b/daiv/sessions/signals.py index 633a92145..a76de17e5 100644 --- a/daiv/sessions/signals.py +++ b/daiv/sessions/signals.py @@ -27,6 +27,10 @@ # operator greps for stays identical across them. LINK_FAILED_PREFIX = "link_failed (agent task will run but its result cannot be captured)" +# ``error_message`` prefix for a run that never reached the broker. Shared with +# ``release_orphan_queued_sessions``, which re-queues continuations matching it. +DISPATCH_FAILED_PREFIX = "dispatch_failed" + @receiver(post_save, sender=settings.AUTH_USER_MODEL) def backfill_session_user(sender: type, instance: Any, created: bool, **kwargs: Any) -> None: @@ -128,7 +132,23 @@ def sync_run_on_task_signal(sender: type, task_result: Any, **kwargs: Any) -> No @receiver(run_finished) def dispatch_next_in_session(sender: type, run: Any, **kwargs: Any) -> None: - """Release queued continuations on this session, one at a time, until one succeeds. + """Release queued continuations on this session via ``release_next_queued``. + + ``skip_dispatch=True`` (passed by re-emits from dispatch-failure paths) suppresses + re-entry while still letting notification receivers fire. + """ + if kwargs.get("skip_dispatch"): + return + + session_id = getattr(run, "session_id", None) + if not session_id: + return + + release_next_queued(session_id) + + +def release_next_queued(session_id: Any) -> None: + """Promote QUEUED Runs on this session to READY and enqueue, one at a time, until one succeeds. Atomic compare-and-swap (``filter(pk=, status=QUEUED).update(status=READY)``) wins the race against concurrent dispatchers reading the same row. A losing race or a @@ -138,30 +158,27 @@ def dispatch_next_in_session(sender: type, run: Any, **kwargs: Any) -> None: by ``MAX_CONSECUTIVE_DISPATCH_FAILURES`` so a broker outage doesn't mass-fail the whole backlog. - ``skip_dispatch=True`` (passed by re-emits from dispatch-failure paths) suppresses - re-entry while still letting notification receivers fire. + Nothing is promoted while ANY sibling is READY/RUNNING: ``run_one_active_per_session`` + only covers the job-pipeline trigger types, so e.g. an active chat run is invisible to + the constraint and a promoted row would just block on the session lock. """ from sessions.models import Run, RunStatus - if kwargs.get("skip_dispatch"): - return - - session_id = getattr(run, "session_id", None) - if not session_id: - return - consecutive_failures = 0 while True: next_q = Run.objects.filter(session_id=session_id, status=RunStatus.QUEUED).order_by("created_at").first() if next_q is None: return + if Run.objects.filter(session_id=session_id, status__in=[RunStatus.READY, RunStatus.RUNNING]).exists(): + return # active sibling; a later release (run_finished, chat turn end, cron) picks this up + try: claimed = Run.objects.filter(pk=next_q.pk, status=RunStatus.QUEUED).update(status=RunStatus.READY) except IntegrityError: # A concurrent insert (e.g. a fresh _submit_one) already created a # READY row on this session; the partial unique constraint blocks us. - logger.debug("dispatch_next_in_session: peer claim on session=%s, backing off", session_id) + logger.debug("release_next_queued: peer claim on session=%s, backing off", session_id) return if claimed != 1: @@ -175,7 +192,7 @@ def dispatch_next_in_session(sender: type, run: Any, **kwargs: Any) -> None: consecutive_failures += 1 if consecutive_failures >= MAX_CONSECUTIVE_DISPATCH_FAILURES: logger.warning( - "dispatch_next_in_session: bailing on session=%s after %d consecutive dispatch failures; " + "release_next_queued: bailing on session=%s after %d consecutive dispatch failures; " "remaining QUEUED siblings left for release_orphan_queued_sessions", session_id, consecutive_failures, @@ -210,7 +227,7 @@ def _enqueue_queued_run(run: Any) -> bool: ) except Exception as err: # noqa: BLE001 logger.exception("dispatch_next_in_session: enqueue failed for run=%s", run.pk) - run.save(update_fields=run.mark_failed("dispatch_failed", err)) + run.save(update_fields=run.mark_failed(DISPATCH_FAILED_PREFIX, err)) emit_run_finished_if_terminal(run, previous_status=RunStatus.READY, skip_dispatch=True) return False @@ -233,6 +250,123 @@ def _enqueue_queued_run(run: Any) -> bool: return True +def render_batch_summary(batch_id: Any, siblings: list) -> str: + """Build the coordinator continuation prompt from denormalized Run fields.""" + from django.urls import reverse + + from sessions.models import RunStatus + + n_ok = sum(1 for r in siblings if r.status == RunStatus.SUCCESSFUL) + n_failed = sum(1 for r in siblings if r.status == RunStatus.FAILED) + lines = [f"The delegated batch {batch_id} has finished ({n_ok} succeeded, {n_failed} failed).", ""] + for r in siblings: + state = "successful" if r.status == RunStatus.SUCCESSFUL else "failed" + checkout = f"{r.repo_id}@{r.ref}" if r.ref else r.repo_id + lines.append(f"## {checkout} ({state})") + lines.append(f"- Session: {reverse('session_detail', kwargs={'thread_id': r.session_id})}") + if r.merge_request_web_url: + lines.append(f"- Merge request: {r.merge_request_web_url}") + summary = (r.result_summary or r.error_message or "").strip() + if summary: + # Quote-indent every reply line so a leg's output cannot float to top level or spoof + # a sibling's section header. + lines.append("- Reply:") + lines.extend(f" > {reply_line}" for reply_line in summary[:500].splitlines()) + elif r.status == RunStatus.FAILED: + # A FAILED leg can arrive with an empty error_message (best-effort terminal save); + # make the gap explicit so the coordinator doesn't read a blank block as a clean no-op. + lines.append("- Reply: (failed with no captured error message; check the leg session)") + lines.append("") + lines.append( + "Compose the consolidated outcome and continue your instructions " + "(e.g. report back where the request came from)." + ) + return "\n".join(lines) + + +@receiver(run_finished) +def resume_coordinator_on_batch_complete(sender: type, run: Any, **kwargs: Any) -> None: + """When every leg of a *delegated* batch is terminal, enqueue one coordinator + continuation run on the parent thread. + + Deliberately ignores ``skip_dispatch`` (unlike ``dispatch_next_in_session``): a last + leg that turns terminal via the dispatch-failure re-emit must still resume the + coordinator. The continuation is always created QUEUED — the insert itself is the + winner election (``run_one_continuation_per_batch``) — then released through the + dispatcher's single promotion path; a busy coordinator (any active run, including + trigger types outside ``run_one_active_per_session``) leaves it QUEUED for a later + release (``run_finished``, chat turn end, or the orphan-release cron). + """ + from sessions.models import Run, RunStatus, Session, SessionOrigin + from sessions.services import acreate_run + + if run.trigger_type != SessionOrigin.DELEGATED_JOB: + return # free in-memory guard: only delegated legs can complete a delegated batch + batch_id = run.batch_id + if not batch_id: + return # continuation runs carry continuation_of_batch_id, not batch_id + + if Run.objects.by_batch(batch_id).exclude(status__in=RunStatus.terminal()).exists(): + return # legs still pending + + if Run.objects.filter(continuation_of_batch_id=batch_id).exists(): + return # already resumed (winner election) + + leg_session = Session.objects.filter(pk=run.session_id).only("parent_thread_id").first() + if leg_session is None or not leg_session.parent_thread_id: + return # not a delegated leg (broadcast batch, or session gone) + parent_thread_id = leg_session.parent_thread_id + + coordinator = Session.objects.select_related("user").filter(thread_id=parent_thread_id).first() + if coordinator is None: + logger.warning( + "resume_coordinator: parent thread %s not found for batch %s; the batch notification is the fallback" + " signal", + parent_thread_id, + batch_id, + ) + return + + siblings = list( + Run.objects.by_batch(batch_id).only( + "status", "repo_id", "ref", "session", "merge_request_web_url", "result_summary", "error_message" + ) + ) + try: + continuation = async_to_sync(acreate_run)( + status=RunStatus.QUEUED, + trigger_type=SessionOrigin.DELEGATED_JOB, + task_result_id=None, + repo_id=coordinator.repo_id, + ref=coordinator.ref, + user=coordinator.user, + prompt=render_batch_summary(batch_id, siblings), + thread_id=parent_thread_id, + sandbox_environment_id=str(coordinator.sandbox_environment_id) + if coordinator.sandbox_environment_id + else None, + continuation_of_batch_id=batch_id, + agent_model=coordinator.agent_model, + agent_thinking_level=coordinator.agent_thinking_level, + ) + except IntegrityError: + return # another worker won run_one_continuation_per_batch + + # QUEUED is the only crash-safe resting state (recoverable by the orphan-release cron); + # promotion + enqueue go through the dispatcher's single release path. + release_next_queued(parent_thread_id) + + continuation.refresh_from_db(fields=["status"]) + if continuation.status == RunStatus.FAILED: + logger.error( + "resume_coordinator: failed to enqueue continuation for batch=%s on parent thread=%s; " + "coordinator will NOT auto-resume until the retry sweep re-queues it — the batch " + "notification is the interim signal", + batch_id, + parent_thread_id, + ) + + @receiver(run_finished, dispatch_uid="agent_sessions.classify_on_run_finished") def classify_on_run_finished(sender: type, run: Any, **kwargs: Any) -> None: """Enqueue post-run classification when a scheduled run reaches a terminal status. diff --git a/daiv/sessions/tasks.py b/daiv/sessions/tasks.py index 6d8d8902c..73a14175f 100644 --- a/daiv/sessions/tasks.py +++ b/daiv/sessions/tasks.py @@ -174,6 +174,24 @@ def sync_stuck_runs_cron_task(): call_command("sync_stuck_runs") +@cron("*/5 * * * *") +@task +@locked_task(key="release-orphan-queued-sessions") +def release_orphan_queued_sessions_cron_task(): + """Release QUEUED Runs stranded on idle sessions (crash/race-recovery backstop). + + The normal release paths are the ``run_finished`` dispatcher and the chat streamer's + turn-end pass; this sweep covers what they can miss: dispatcher bail-outs after + consecutive enqueue failures (broker outage) and rows whose release signal was lost. + The wrapped command also re-queues delegated-batch continuations that FAILED before + ever starting, so a transient dispatch failure cannot permanently strand a coordinator. + + ``locked_task`` (non-blocking) skips this tick if a prior run still holds the lock, so a + pass that overruns the interval is never double-dispatched. + """ + call_command("release_orphan_queued_sessions") + + @cron("*/15 * * * *") @task @locked_task(key="reclassify-missing-envelopes") diff --git a/docs/customization/repository-config.md b/docs/customization/repository-config.md index 634ed7743..f99687e3c 100644 --- a/docs/customization/repository-config.md +++ b/docs/customization/repository-config.md @@ -31,6 +31,9 @@ pull_request_assistant: slash_commands: enabled: true +orchestration: + enabled: true + # Model overrides models: agent: @@ -96,6 +99,9 @@ pull_request_assistant: slash_commands: enabled: false + +orchestration: + enabled: false ``` | Section | Option | Default | Description | @@ -103,6 +109,7 @@ slash_commands: | `issue_addressing` | `enabled` | `true` | [Issue Addressing](../features/issue-addressing.md) | | `pull_request_assistant` | `enabled` | `true` | [Pull Request Assistant](../features/pull-request-assistant.md) | | `slash_commands` | `enabled` | `true` | [Slash Commands & Skills](../features/slash-commands.md) | +| `orchestration` | `enabled` | `true` | [Orchestration](../features/orchestration.md) — binds the `delegate_jobs` tool | ## Branch naming and commit conventions diff --git a/docs/features/jobs-api.md b/docs/features/jobs-api.md index e0586a964..416d2601d 100644 --- a/docs/features/jobs-api.md +++ b/docs/features/jobs-api.md @@ -69,8 +69,8 @@ POST /api/jobs | Field | Type | Required | Description | |------------------------|------------------|----------|-------------| -| `repos` | array of objects | yes | 1–20 repositories to run against. Each item: `{ "repo_id": "group/project", "ref": "branch-or-sha" }` — `ref` is optional and defaults to the repository's default branch. | -| `prompt` | string | yes | The prompt to send to the agent. The same prompt runs as an independent job against each repository in `repos`. | +| `repos` | array of objects | yes | 1–20 repositories to run against. Each item: `{ "repo_id": "group/project", "ref": "branch-or-sha", "prompt": "..." }` — `ref` and `prompt` are optional. `ref` defaults to the repository's default branch; `repos[].prompt` overrides the batch-level `prompt` for that repository. | +| `prompt` | string | yes | The prompt to send to the agent. The same prompt runs as an independent job against each repository in `repos`, unless a per-repo override is provided via `repos[].prompt`. | | `agent_model` | string | no | Override the model used for this batch. Invalid model / thinking-level combinations are rejected with `400`. | | `agent_thinking_level` | string | no | Override the agent's reasoning effort. One of `minimal`, `low`, `medium`, `high`, `xhigh`. Invalid combinations are rejected with `400`. | | `notify_on` | string | no | Override the user's notification preference for this batch. One of `never`, `always`, `on_success`, `on_failure`. | diff --git a/docs/features/mcp-endpoint.md b/docs/features/mcp-endpoint.md index d57bea00b..e6c02f6d5 100644 --- a/docs/features/mcp-endpoint.md +++ b/docs/features/mcp-endpoint.md @@ -87,7 +87,7 @@ url = "https://daiv.example.com/mcp/" The paginated listing tools (`list_jobs`, `list_scheduled_jobs`, `list_environments`) share one pagination contract: pass an optional `limit` and `cursor`, and read back `{ "": [...], "next_cursor": }`. To page, call again with `cursor` set to the previous response's `next_cursor` until it comes back `null`; a cursor encodes only sort position, so reuse it with the **same** filters. `list_repositories` is also served from the database but never paginates — it returns the same shape for consistency (narrow with `search`/`topics` instead). -`submit_job` takes a `repos` list (1–20 entries) and a single `prompt` that runs as an independent job against each repository. Each entry is `{repo_id, ref}`, where `ref` is the starting branch or commit the agent reads from — it is optional and defaults to the repository's default branch. The response includes a `batch_id`, a `jobs` list (one entry per submitted job, each with its `job_id`, `repo_id`, `ref`, `thread_id`, and `status`), and a `failed` list for repositories that could not be enqueued. +`submit_job` takes a `repos` list (1–20 entries) and a single `prompt` that runs as an independent job against each repository. Each entry is `{repo_id, ref, prompt}`, where `ref` is the starting branch or commit the agent reads from (optional, defaults to the repository's default branch), and `repos[].prompt` is an optional per-repo instruction that overrides the batch-level `prompt` for that repository. The response includes a `batch_id`, a `jobs` list (one entry per submitted job, each with its `job_id`, `repo_id`, `ref`, `thread_id`, and `status`), and a `failed` list for repositories that could not be enqueued. `submit_job` also accepts these optional parameters: diff --git a/docs/features/orchestration.md b/docs/features/orchestration.md new file mode 100644 index 000000000..3eebeaf23 --- /dev/null +++ b/docs/features/orchestration.md @@ -0,0 +1,128 @@ +# Orchestration + +Orchestration lets a DAIV agent fan work out to other repositories in a single step. When a task spans multiple codebases — shared libraries, microservices, a monorepo alongside its satellite packages — the agent can delegate sub-tasks in parallel rather than tackling each repository one after another. + +This is useful when you want to: + +- **Coordinate multi-repo changes** — e.g., bump a shared library version across every service that depends on it +- **Triage and route tickets automatically** — e.g., read an issue queue and dispatch each ticket to the right repository +- **Run parallel investigations** — e.g., check all affected services for a security advisory +- **Build a coordination-repo workflow** — one repository acts as the orchestrator; all the real work lands in the target repos + +## How it works + +The agent has access to a `delegate_jobs` tool, enabled by default on every repository. The tool accepts a `goal` and a list of target repositories (each with its own tailored `prompt`) and submits an independent agent run for every target — in parallel, via the same task backend used by the Jobs API and Scheduled Jobs. + +Each delegated run executes with the per-repo configuration, skills, and sandbox of its own repository. When all delegated runs finish, the originating agent resumes and receives a rollup summary of the results — what each run produced, whether it succeeded, and any merge requests that were created. + +```mermaid +sequenceDiagram + participant Coord as Coordinator agent + participant Tool as delegate_jobs tool + participant RunA as Agent run — repo A + participant RunB as Agent run — repo B + participant RunC as Agent run — repo C + + Coord->>Tool: delegate_jobs(goal, [repo_A, repo_B, repo_C]) + Tool->>RunA: submit (prompt, repo_A) + Tool->>RunB: submit (prompt, repo_B) + Tool->>RunC: submit (prompt, repo_C) + RunA-->>Tool: done + RunB-->>Tool: done + RunC-->>Tool: done + Tool-->>Coord: rollup (results, MR URLs) +``` + +## Enabling and disabling orchestration + +Orchestration is **enabled by default** — the `delegate_jobs` tool is bound to the agent on every repository. To turn it off for a repository, set the following in its `.daiv.yml`: + +```yaml +orchestration: + enabled: false +``` + +When disabled, the `delegate_jobs` tool is not bound to the agent, and any attempt to delegate results in an error. + +!!! note + A target repository needs no configuration to *receive* delegated work — it just runs a normal agent job. Disabling orchestration on a repository only stops that repository's own agent from *initiating* delegation. + +## Limits + +| Limit | Value | Notes | +|-------|-------|-------| +| **Width** — targets per `delegate_jobs` call | 10 | The tool rejects a call with more than 10 `targets`. Split larger fan-outs across multiple calls. | +| **Depth** — maximum delegation chain | 2 | A delegated run cannot itself delegate beyond this depth. Setting `MAX_SPAWN_DEPTH=2` means: coordinator (depth 0) → delegated leg (depth 1) → leaf leg (depth 2) → no further delegation. | + +!!! warning + Depth is enforced at submission time (and pinned by a database constraint on the session). A delegated run that tries to call `delegate_jobs` when it is already at the maximum depth will receive an error and should handle it gracefully in its prompt. Because orchestration is on by default, a delegated leg may itself delegate *further* (up to the depth cap) unless its repository sets `orchestration.enabled: false`. + +## Per-target prompts + +Every `delegate_jobs` target **requires** its own `prompt`. A delegated leg runs in isolation and cannot see the coordinator's ticket or context, so each leg needs explicit, self-contained instructions — include the ticket context that leg needs and the no-change convention ("if this repository is unaffected, reply saying so and make no changes"): + +```json +{ + "goal": "Apply the CVE-2026-12345 patch from the advisory", + "targets": [ + { "repo_id": "mygroup/service-auth", "ref": "main", "prompt": "Bump the vulnerable auth library to 3.2.1; if this repo doesn't depend on it, reply saying so and make no changes" }, + { "repo_id": "mygroup/service-api", "ref": "main", "prompt": "Update the auth dependency to 3.2.1 and run the integration tests" }, + { "repo_id": "mygroup/service-web", "ref": "main", "prompt": "Rebuild against auth 3.2.1 and refresh the lockfile; reply with no changes if unaffected" } + ] +} +``` + +The batch-level `goal` is a one-line objective used only to **title the batch** in the session view — it is never sent to a leg as a prompt. + +!!! note + A per-repo prompt that is *optional* and falls back to a shared batch prompt is a feature of the [Jobs API](jobs-api.md) and [MCP `submit_job`](mcp-endpoint.md), where `repos[].prompt` may be omitted. The `delegate_jobs` tool is deliberately stricter: because a leg cannot see the coordinator's context, every target must be told explicitly what to do. + +## Ticket-triage recipe + +A common pattern is a **coordination repository** that holds no production code but acts as a routing hub for an external issue queue. The setup has three parts: + +### 1. Attach a ticketing MCP server + +Configure the coordination repository with an MCP tool integration that points at your ticketing system (e.g. Request Tracker, Jira, Linear). See [MCP Tools](../customization/mcp-tools.md) for how to wire up a per-repo MCP server. + +### 2. Put routing knowledge in `.agents/AGENTS.md` + +The coordination repository's `.agents/AGENTS.md` is the right place to describe how to map tickets to repositories: + +```markdown +# Routing rules + +- Tickets tagged `backend` or assigned to queue `Platform` → mygroup/service-api +- Tickets tagged `frontend` → mygroup/service-web +- Tickets tagged `auth` → mygroup/service-auth +- Unclassified tickets → skip, add comment "needs triage label" +``` + +The agent reads this file at the start of every run. Keep the routing rules up to date as your team's structure changes. + +### 3. Trigger on a schedule (optional) + +Use a [Scheduled Job](scheduled-jobs.md) on the coordination repository to poll the ticket queue at a regular interval — e.g., every hour — and dispatch the results automatically, without any manual prompt: + +``` +Scheduled prompt: +Fetch all open tickets from the RT queue that have not yet been dispatched. +Follow the routing rules in AGENTS.md to pick each ticket's target repository, then +delegate_jobs with one target per repository — when several tickets route to the same +repository, combine them into that target's prompt (duplicate targets are rejected). +Add a comment to each ticket with the resulting merge request URL. +``` + +Combine this with the [Jobs API](jobs-api.md) to trigger an immediate dispatch from a webhook (e.g., on ticket-create) instead of — or alongside — the scheduled poller. + +## Session view + +Each `delegate_jobs` call produces a batch of child runs visible in the [Sessions](sessions.md) list. The coordinator run links to the batch; each child run shows its own transcript, run timeline, and any merge requests it produced. When all children finish, the coordinator resumes and the rollup appears in the coordinator's session transcript. + +## See also + +- [Jobs API](jobs-api.md) — programmatic job submission, including `thread_id` for continuation +- [Scheduled Jobs](scheduled-jobs.md) — recurring agent runs that can seed an orchestration workflow +- [MCP Endpoint](mcp-endpoint.md) — expose DAIV jobs to AI coding assistants via MCP +- [Repository Config](../customization/repository-config.md) — full reference for `.daiv.yml` +- [Agent Skills](../customization/agent-skills.md) — per-repo skills loaded into agent runs diff --git a/docs/index.md b/docs/index.md index bf6842afc..9bd0cf7ec 100644 --- a/docs/index.md +++ b/docs/index.md @@ -31,6 +31,7 @@ DAIV integrates directly with **GitLab** and **GitHub** through webhooks — no - **[MCP Endpoint](features/mcp-endpoint.md)** — connect Claude Code, Cursor, or Codex CLI over the [Model Context Protocol](https://modelcontextprotocol.io/) and delegate tasks without leaving your editor. - **[Jobs API](features/jobs-api.md)** — trigger agents programmatically from CI, scripts, or other tools, then poll for the result. +- **[Orchestration](features/orchestration.md)** — fan work out across repositories: one coordinator agent delegates sub-tasks to other repos in parallel and resumes with a rollup of the results. ### From the dashboard diff --git a/docs/reference/env-variables.md b/docs/reference/env-variables.md index 2d63dfc5b..d1f4aeaeb 100644 --- a/docs/reference/env-variables.md +++ b/docs/reference/env-variables.md @@ -345,6 +345,15 @@ The main agent used for issue addressing, pull request assistance, and all inter | `DAIV_AGENT_EXPLORE_FALLBACK_MODEL_NAME` | Fallback model if the explore model fails | `gpt-5-4-mini` | | `DAIV_AGENT_CUSTOM_SKILLS_PATH` | Path to custom global skills directory. Set to `None` to disable. | `~/data/skills` | +### Model Requests + +These bound every LLM API call — the main agent, subagents, and the model-backed tasks (titling, memory, commit/PR writer, web fetch) — regardless of provider. Without them a stalled provider can block a run indefinitely. + +| Variable | Description | Default | +|-------------------------------------|--------------------------------------------------------------------------|---------| +| `DAIV_MODEL_REQUEST_TIMEOUT_SECONDS` | Per-request timeout for LLM API calls, in seconds. Minimum `1`. | `600` | +| `DAIV_MODEL_MAX_RETRIES` | Retries for a failed LLM API call before the error propagates (`0` disables retries) | `1` | + ### Jobs API The [Jobs API](../features/jobs-api.md) allows programmatic agent execution. diff --git a/docs/reference/site-configuration.md b/docs/reference/site-configuration.md index 337e20044..128f94081 100644 --- a/docs/reference/site-configuration.md +++ b/docs/reference/site-configuration.md @@ -19,7 +19,7 @@ The configuration is split into the following groups, organized by category. | Group | Category | What it configures | |-------|----------|--------------------| -| **Agent** | AI tasks | Primary, fallback, `daiv-max`, and explore models; their thinking levels; the agent recursion limit; and whether to suggest a context file (e.g. `AGENTS.md`) on new merge requests. | +| **Agent** | AI tasks | Primary, fallback, `daiv-max`, and explore models; their thinking levels; the agent recursion limit; the per-request LLM timeout and max retries (applied to every provider and model-backed task); and whether to suggest a context file (e.g. `AGENTS.md`) on new merge requests. | | **Commit & PR Writer** | AI tasks | Primary and fallback models that generate commit messages and pull/merge request descriptions from diffs. | | **Titling** | AI tasks | Primary and fallback models that generate session and run titles from prompts. | | **Providers** | Models | LLM provider records — slug, wire protocol, base URL, API key, and extra headers — that every model field draws from. | diff --git a/mkdocs.yml b/mkdocs.yml index 089b81244..964d6e094 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -82,6 +82,7 @@ plugins: - features/sessions.md: Unified sessions — chat workspace, run history, and live updates - features/jobs-api.md: REST API for programmatic agent triggering - features/scheduled-jobs.md: Recurring agent runs on a cron schedule + - features/orchestration.md: Fan-out agent runs across multiple repositories from a single coordinator - features/mcp-endpoint.md: MCP server endpoint for AI coding assistants - features/notifications.md: Job, batch, and schedule completion notifications - features/merge-metrics.md: Code merge analytics and DAIV contribution tracking @@ -144,6 +145,7 @@ nav: - Sessions: features/sessions.md - Jobs API: features/jobs-api.md - Scheduled Jobs: features/scheduled-jobs.md + - Orchestration: features/orchestration.md - MCP Endpoint: features/mcp-endpoint.md - Notifications: features/notifications.md - Merge Metrics: features/merge-metrics.md diff --git a/tests/unit_tests/automation/agent/middlewares/test_web_fetch.py b/tests/unit_tests/automation/agent/middlewares/test_web_fetch.py index e02e10f1f..18b4f24dc 100644 --- a/tests/unit_tests/automation/agent/middlewares/test_web_fetch.py +++ b/tests/unit_tests/automation/agent/middlewares/test_web_fetch.py @@ -168,6 +168,11 @@ async def test_caches_full_response_by_url_and_prompt(httpx_mock): assert result1 == "ANSWER" assert result2 == "ANSWER" assert httpx_mock.get_requests() + # Second call is served from cache, so the summariser model is built exactly once, + # and it must carry the tighter web_fetch budget rather than the global default. + mock_base_agent.get_model.assert_called_once_with( + model="openrouter:openai/gpt-4.1-mini", timeout=web_fetch_module.WEB_FETCH_MODEL_TIMEOUT_SECONDS + ) async def test_cache_key_changes_with_prompt(httpx_mock): @@ -339,3 +344,32 @@ async def ainvoke(self, _messages): result = await web_fetch_module.web_fetch_tool.ainvoke({"url": "https://example.com", "prompt": "x"}) assert result == "Model processing failed (Boom). Contents of https://example.com:\nCONTENT" + + +async def test_model_build_failure_logs_error_and_still_returns_content(httpx_mock, caplog): + """A permanent misconfiguration (get_model raising, e.g. missing API key) must be logged at + ERROR (Sentry-visible) for operators, yet still return the raw page content — flagged so the + agent knows it is unsummarised — because the agent can reason over it and cannot fix config.""" + httpx_mock.add_response( + url="https://example.com", + status_code=200, + headers={"content-type": "text/html"}, + text="CONTENT", + ) + with ( + patch.object(web_fetch_module, "BaseAgent") as mock_base_agent, + patch.object(web_fetch_module, "site_settings") as mock_site_settings, + patch.object(web_fetch_module, "automation_env_settings") as mock_env_settings, + ): + mock_site_settings.web_fetch_timeout_seconds = 1 + mock_env_settings.WEB_FETCH_PROXY_URL = None + mock_site_settings.web_fetch_max_content_chars = 999_999 + mock_site_settings.web_fetch_model_name = "openrouter:openai/gpt-4.1-mini" + mock_base_agent.get_model.side_effect = RuntimeError("Provider 'openrouter' has no API key configured.") + + with caplog.at_level("ERROR", logger="daiv.tools"): + result = await web_fetch_module.web_fetch_tool.ainvoke({"url": "https://example.com", "prompt": "x"}) + + assert "CONTENT" in result # raw content is preserved, not thrown away + assert "summariser unavailable" in result # but clearly flagged as unsummarised + assert any(rec.levelname == "ERROR" for rec in caplog.records) # operators still get the signal diff --git a/tests/unit_tests/automation/agent/test_base.py b/tests/unit_tests/automation/agent/test_base.py index 91fd538c1..35b0ee821 100644 --- a/tests/unit_tests/automation/agent/test_base.py +++ b/tests/unit_tests/automation/agent/test_base.py @@ -5,7 +5,7 @@ from langchain.chat_models import BaseChatModel from langchain_core.runnables import Runnable -from automation.agent.base import BaseAgent, ResolvedProvider, parse_model_spec +from automation.agent.base import BaseAgent, ResolvedProvider, _resolve_request_timeout, parse_model_spec from automation.agent.chat_models import OPENROUTER_BASE_URL, ChatOpenRouter from core.models import Provider, ProviderType, ThinkingLevelChoices @@ -455,3 +455,102 @@ def test_adaptive_generation_openrouter_thinking_without_temperature(self, model assert kw["extra_body"]["reasoning"]["enabled"] is True assert kw["extra_body"]["reasoning"]["effort"] == ThinkingLevelChoices.MEDIUM assert "temperature" not in kw + + @pytest.mark.parametrize( + ("slug", "api_key", "model_spec", "expects_httpx_timeout", "expected_max_retries"), + [ + # OpenAI-family accepts an httpx.Timeout (connect fuse); Anthropic/Google a bare float. + # google-genai's max_retries counts total attempts, so "1 retry" reaches it as 2. + ("openai", "sk-o", "openai:gpt-5.4", True, 1), + ("openrouter", "sk-or", "openrouter:z-ai/glm-5", True, 1), + ("anthropic", "sk-a", "anthropic:claude-sonnet-4-6", False, 1), + ("google_genai", "sk-g", "google_genai:gemini-2.5-pro", False, 2), + ], + ) + def test_timeout_and_retries_applied_per_provider( + self, slug, api_key, model_spec, expects_httpx_timeout, expected_max_retries + ): + """Every provider gets the bounded default (600s / 1 retry) in the timeout and retry + shapes its langchain integration accepts. The per-shape edge cases (connect fuse, + scalar-vs-Timeout) are covered directly in ``TestResolveRequestTimeout``.""" + import httpx + + self._enable_seed(slug, api_key) + kw = BaseAgent.get_model_kwargs(resolved=parse_model_spec(model_spec)) + assert kw["max_retries"] == expected_max_retries + if expects_httpx_timeout: + assert isinstance(kw["timeout"], httpx.Timeout) + assert kw["timeout"].read == 600.0 + assert kw["timeout"].connect == 10.0 + else: + assert kw["timeout"] == 600.0 + assert not isinstance(kw["timeout"], httpx.Timeout) + + def test_caller_timeout_and_retries_take_precedence(self): + """An explicit per-call budget (e.g. the tighter web_fetch timeout) is never + overridden by the global defaults, and stays verbatim rather than being wrapped.""" + import httpx + + self._enable_seed("openai", "sk-o") + kw = BaseAgent.get_model_kwargs(resolved=parse_model_spec("openai:gpt-5.4"), timeout=45, max_retries=5) + assert kw["timeout"] == 45 + assert kw["max_retries"] == 5 + # A caller override is passed through raw — even for OpenAI, it is not re-wrapped in + # httpx.Timeout (this is the web_fetch path, get_model(..., timeout=60)). + assert not isinstance(kw["timeout"], httpx.Timeout) + + def test_insecure_http_clients_inherit_resolved_timeout(self): + """The custom verify=False clients must carry the resolved timeout instead of httpx's + 5s default, which would otherwise throttle real generations.""" + import httpx + + Provider.objects.create( + slug="insec_timeout", + display_name="Insecure", + provider_type=ProviderType.OPENAI, + api_key="sk-x", + base_url="https://internal.example.test/v1", + verify_ssl=False, + ) + kw = BaseAgent.get_model_kwargs(resolved=parse_model_spec("insec_timeout:gpt-5.4")) + assert isinstance(kw["timeout"], httpx.Timeout) + assert kw["http_client"].timeout.read == 600.0 + assert kw["http_async_client"].timeout.read == 600.0 + + def test_explicit_none_override_does_not_unbind(self): + """A caller passing ``timeout=None`` / ``max_retries=None`` must NOT be able to silently + restore the unbounded default — an explicit None is treated as unset.""" + import httpx + + self._enable_seed("openai", "sk-o") + kw = BaseAgent.get_model_kwargs(resolved=parse_model_spec("openai:gpt-5.4"), timeout=None, max_retries=None) + assert isinstance(kw["timeout"], httpx.Timeout) + assert kw["timeout"].read == 600.0 + assert kw["max_retries"] == 1 + + +class TestResolveRequestTimeout: + """Direct tests for the per-provider timeout shaping — exercises the ``connect`` fuse on the + sub-10s branch, which the default-driven get_model_kwargs tests never reach.""" + + def test_openai_short_timeout_caps_connect_to_total(self): + """When the total budget is below 10s, connect must not exceed it (the reason ``min`` exists).""" + import httpx + + t = _resolve_request_timeout(ProviderType.OPENAI, 5) + assert isinstance(t, httpx.Timeout) + assert t.read == 5.0 + assert t.connect == 5.0 + + def test_openai_long_timeout_uses_10s_connect_fuse(self): + import httpx + + t = _resolve_request_timeout(ProviderType.OPENROUTER, 120) + assert isinstance(t, httpx.Timeout) + assert t.read == 120.0 + assert t.connect == 10.0 + + def test_anthropic_and_google_return_bare_float(self): + """Neither langchain integration accepts an httpx.Timeout — the value stays a scalar.""" + assert _resolve_request_timeout(ProviderType.ANTHROPIC, 120) == 120.0 + assert _resolve_request_timeout(ProviderType.GOOGLE_GENAI, 5) == 5.0 diff --git a/tests/unit_tests/automation/agent/test_delegate_jobs_tool.py b/tests/unit_tests/automation/agent/test_delegate_jobs_tool.py new file mode 100644 index 000000000..d4f5926fa --- /dev/null +++ b/tests/unit_tests/automation/agent/test_delegate_jobs_tool.py @@ -0,0 +1,215 @@ +import json +from contextlib import contextmanager +from unittest.mock import AsyncMock, patch + +import pytest +from sessions.models import Session, SessionOrigin + +from automation.agent.middlewares.delegate_jobs import delegate_jobs_tool + +pytestmark = pytest.mark.django_db(transaction=True) + +CONFIG = {"configurable": {"thread_id": "coord-thread"}} + + +async def _invoke(goal, targets): + return json.loads(await delegate_jobs_tool.ainvoke({"goal": goal, "targets": targets}, config=CONFIG)) + + +def _fake_result(repo_id="g/a", session_id="leg-1", batch_id="batch-1"): + run = type("R", (), {"repo_id": repo_id, "ref": "", "session_id": session_id})() + return type("B", (), {"batch_id": batch_id, "runs": [run], "failed": []})() + + +async def _catalog(repo_id, default_branch): + from django.utils import timezone + + from codebase.conf import settings as codebase_settings + from codebase.models import RepositoryCatalog + + await RepositoryCatalog.objects.acreate( + provider=codebase_settings.CLIENT.value, slug=repo_id, default_branch=default_branch, synced_at=timezone.now() + ) + + +@contextmanager +def _patched_delegate(submit): + """Patch the tool's auth + env-resolution + batch-submit collaborators. + + ``submit`` is the AsyncMock used for ``asubmit_batch_runs`` (a return_value for the happy path or + a side_effect to simulate a raise); it is yielded so callers can assert on its call args. + """ + with ( + patch("automation.agent.middlewares.delegate_jobs.aassert_can_run", new=AsyncMock(return_value=None)), + patch( + "automation.agent.middlewares.delegate_jobs.aresolve_repo_envs", + new=AsyncMock(side_effect=lambda **kw: kw["repos"]), + ), + patch("automation.agent.middlewares.delegate_jobs.asubmit_batch_runs", new=submit), + ): + yield submit + + +async def test_refuses_when_session_has_no_user(django_user_model): + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=None) + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "p"}]) + assert "no user" in out["error"].lower() + + +async def test_refuses_at_depth_cap(django_user_model): + user = await django_user_model.objects.acreate(username="u1") + await Session.objects.acreate( + thread_id="coord-thread", origin=SessionOrigin.DELEGATED_JOB, repo_id="g/coord", user=user, spawn_depth=2 + ) + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "p"}]) + assert "depth" in out["error"].lower() + + +async def test_denied_targets_reported_inline(django_user_model): + user = await django_user_model.objects.acreate(username="u2") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + + from codebase.authorization import RepositoryAccessDenied + + with patch( + "automation.agent.middlewares.delegate_jobs.aassert_can_run", + new=AsyncMock(side_effect=RepositoryAccessDenied(["g/a", "g/b"])), + ): + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "p"}, {"repo_id": "g/b", "prompt": "q"}]) + assert out["batch_id"] is None + assert {f["repo_id"] for f in out["failed"]} == {"g/a", "g/b"} + # Nothing was started, so the model must be told not to wait for a resume. + assert "handle the failures now" in out["note"] + + +async def test_success_path_submits_allowed_targets(django_user_model): + user = await django_user_model.objects.acreate(username="u3") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + + with _patched_delegate(AsyncMock(return_value=_fake_result())) as m_submit: + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "do X"}]) + + assert out["batch_id"] == "batch-1" + assert out["delegated"][0]["repo_id"] == "g/a" + assert out["delegated"][0]["session_url"] == "/dashboard/sessions/leg-1/" + # parent_thread_id + spawn_depth were passed through + kwargs = m_submit.call_args.kwargs + assert kwargs["parent_thread_id"] == "coord-thread" + assert kwargs["spawn_depth"] == 1 + assert kwargs["trigger_type"] == SessionOrigin.DELEGATED_JOB + assert "note" not in out + + +async def test_refuses_empty_targets(django_user_model): + user = await django_user_model.objects.acreate(username="u-empty") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + out = await _invoke("goal", []) + assert "at least one target" in out["error"].lower() + + +async def test_refuses_more_than_max_targets(django_user_model): + from sessions.services import MAX_DELEGATED_TARGETS + + user = await django_user_model.objects.acreate(username="u-many") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + targets = [{"repo_id": f"g/r{i}", "prompt": "p"} for i in range(MAX_DELEGATED_TARGETS + 1)] + out = await _invoke("goal", targets) + assert "at most" in out["error"].lower() + + +async def test_refuses_duplicate_target(django_user_model): + user = await django_user_model.objects.acreate(username="u-dup") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + # Same repo, both with the default (omitted) ref → collide on ("g/a", ""). + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "p"}, {"repo_id": "g/a", "prompt": "q"}]) + assert "duplicate target" in out["error"].lower() + + +async def test_refuses_self_delegation(django_user_model): + """Delegating to the coordinator's own repo+ref is refused and steered to subagents (`task`).""" + user = await django_user_model.objects.acreate(username="u-self") + await Session.objects.acreate( + thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", ref="main", user=user + ) + out = await _invoke("goal", [{"repo_id": "g/coord", "ref": "main", "prompt": "p"}]) + assert "g/coord" in out["error"] + assert "task" in out["error"].lower() + + +async def test_self_delegation_fails_whole_call_without_submitting(django_user_model): + """A self-target (same repo+ref) mixed with valid targets fails the entire call — nothing submitted.""" + user = await django_user_model.objects.acreate(username="u-self-mix") + await Session.objects.acreate( + thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", ref="main", user=user + ) + with _patched_delegate(AsyncMock(return_value=_fake_result())) as m_submit: + out = await _invoke( + "goal", [{"repo_id": "g/other", "prompt": "p"}, {"repo_id": "g/coord", "ref": "main", "prompt": "q"}] + ) + assert "task" in out["error"].lower() + m_submit.assert_not_called() + + +async def test_refuses_self_delegation_when_target_names_default_branch(django_user_model): + """An omitted coordinator ref and the default branch's explicit name are the same checkout.""" + user = await django_user_model.objects.acreate(username="u-self-alias") + await Session.objects.acreate( + thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", ref="", user=user + ) + await _catalog("g/coord", "main") + out = await _invoke("goal", [{"repo_id": "g/coord", "ref": "main", "prompt": "p"}]) + assert "task" in out["error"].lower() + + +async def test_refuses_duplicate_target_via_default_branch_alias(django_user_model): + """(repo, omitted ref) and (repo, explicit default-branch name) are one checkout, not two targets.""" + user = await django_user_model.objects.acreate(username="u-dup-alias") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + await _catalog("g/a", "main") + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "p"}, {"repo_id": "g/a", "ref": "main", "prompt": "q"}]) + assert "duplicate target" in out["error"].lower() + + +async def test_allows_same_repo_different_ref(django_user_model): + """Same repo on a *different* ref is a distinct checkout, not self-delegation — it delegates.""" + user = await django_user_model.objects.acreate(username="u-diff-ref") + await Session.objects.acreate( + thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", ref="main", user=user + ) + with _patched_delegate(AsyncMock(return_value=_fake_result())) as m_submit: + out = await _invoke("goal", [{"repo_id": "g/coord", "ref": "feature-x", "prompt": "p"}]) + assert out["batch_id"] == "batch-1" + m_submit.assert_called_once() + + +async def test_allows_just_under_depth_cap(django_user_model): + """A coordinator at spawn_depth=1 (one below the cap) delegates, stamping legs at depth 2.""" + user = await django_user_model.objects.acreate(username="u-boundary") + await Session.objects.acreate( + thread_id="coord-thread", origin=SessionOrigin.DELEGATED_JOB, repo_id="g/coord", user=user, spawn_depth=1 + ) + + with _patched_delegate(AsyncMock(return_value=_fake_result())) as m_submit: + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "do X"}]) + + assert out["batch_id"] == "batch-1" + assert m_submit.call_args.kwargs["spawn_depth"] == 2 + + +async def test_submission_failure_returns_json_error(django_user_model): + """A raise from env resolution / batch submit is reported as a JSON error, not propagated.""" + user = await django_user_model.objects.acreate(username="u-boom") + await Session.objects.acreate(thread_id="coord-thread", origin=SessionOrigin.MCP_JOB, repo_id="g/coord", user=user) + + with _patched_delegate(AsyncMock(side_effect=RuntimeError("db exploded"))): + out = await _invoke("goal", [{"repo_id": "g/a", "prompt": "do X"}]) + + assert "error" in out + assert "submission failed" in out["error"].lower() + + +def test_middleware_exposes_the_tool(): + from automation.agent.middlewares.delegate_jobs import DELEGATE_JOBS_NAME, DelegateJobsMiddleware + + mw = DelegateJobsMiddleware() + assert [t.name for t in mw.tools] == [DELEGATE_JOBS_NAME] diff --git a/tests/unit_tests/chat/api/test_streaming.py b/tests/unit_tests/chat/api/test_streaming.py index 895cd36e0..4e132e806 100644 --- a/tests/unit_tests/chat/api/test_streaming.py +++ b/tests/unit_tests/chat/api/test_streaming.py @@ -111,6 +111,26 @@ async def _capture_release(thread_id, run_id): assert release_calls == [("t-stream", "r-1")] +@pytest.mark.django_db(transaction=True) +async def test_events_releases_queued_runs_at_turn_end(): + """Chat runs emit no run_finished, so the turn-end pass must release QUEUED rows + (e.g. a delegated-batch continuation parked while this turn held the session lock).""" + with ( + patch("chat.api.streaming.open_checkpointer", _mock_ctx), + patch("chat.api.streaming.set_runtime_ctx", _mock_ctx), + patch("chat.api.streaming.create_daiv_agent", new=AsyncMock()), + patch("chat.api.streaming.RuntimeContextLangGraphAGUIAgent", return_value=_mock_agent([])), + patch("chat.api.streaming.SessionLock.release", new=AsyncMock()), + patch("chat.api.streaming.SessionLock.heartbeat", new=AsyncMock()), + patch("sessions.signals.release_next_queued") as m_release, + ): + streamer = _streamer() + async for _ in streamer.events(): + pass + + m_release.assert_called_once_with("t-stream") + + @pytest.mark.django_db(transaction=True) async def test_events_captures_latest_merge_request_when_multiple_snapshots(): """Multiple snapshots arrive — last one wins. Regression for diff --git a/tests/unit_tests/codebase/test_repo_config.py b/tests/unit_tests/codebase/test_repo_config.py index fa06a541a..10e93b380 100644 --- a/tests/unit_tests/codebase/test_repo_config.py +++ b/tests/unit_tests/codebase/test_repo_config.py @@ -58,3 +58,17 @@ def test_memory_section_can_be_disabled(): config = RepositoryConfig(**{"memory": {"enabled": False}}) assert config.memory.enabled is False + + +def test_orchestration_defaults_on(): + from codebase.repo_config import RepositoryConfig + + cfg = RepositoryConfig() + assert cfg.orchestration.enabled is True + + +def test_orchestration_can_be_disabled_via_yaml(): + from codebase.repo_config import RepositoryConfig + + cfg = RepositoryConfig.model_validate({"orchestration": {"enabled": False}}) + assert cfg.orchestration.enabled is False diff --git a/tests/unit_tests/jobs/api/test_views.py b/tests/unit_tests/jobs/api/test_views.py index 7ab6e9d8a..37bb542e8 100644 --- a/tests/unit_tests/jobs/api/test_views.py +++ b/tests/unit_tests/jobs/api/test_views.py @@ -506,3 +506,31 @@ async def test_submit_job_denied_repo_returns_opaque_404(authenticated_client: T assert response.status_code == 404 assert response.json()["detail"] == "Repository not found or not accessible." + + +@pytest.mark.django_db(transaction=True) +async def test_rest_per_repo_prompt_passthrough(authenticated_client: TestAsyncClient): + """Per-repo ``prompt`` on ``RepoSubmitItem`` reaches the ``RepoTarget`` passed to ``asubmit_batch_runs``.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from sessions.services import BatchSubmitFailure, BatchSubmitResult + + # Report the single repo as failed so the view's runs_iter loop is never entered, + # keeping the test focused on the target-build path. + failure = BatchSubmitFailure(repo_id="group/project", ref="", error="dry-run") + batch_result = MagicMock(spec=BatchSubmitResult) + batch_result.batch_id = "00000000-0000-0000-0000-000000000001" + batch_result.runs = [] + batch_result.failed = [failure] + + submit = AsyncMock(return_value=batch_result) + with patch("jobs.api.views.asubmit_batch_runs", new=submit): + response = await authenticated_client.post( + "/jobs", + json={"repos": [{"repo_id": "group/project", "prompt": "per-repo override"}], "prompt": "batch-level"}, + ) + + assert response.status_code == 202 + targets = submit.call_args.kwargs["repos"] + assert len(targets) == 1 + assert targets[0].prompt == "per-repo override" diff --git a/tests/unit_tests/mcp_server/test_server_jobs.py b/tests/unit_tests/mcp_server/test_server_jobs.py index 3dc469485..4be315e9c 100644 --- a/tests/unit_tests/mcp_server/test_server_jobs.py +++ b/tests/unit_tests/mcp_server/test_server_jobs.py @@ -272,3 +272,21 @@ async def test_submit_job_rate_limited(): result = await submit_job(repos=[{"repo_id": "a/b", "ref": None}], prompt="x") assert "Rate limit" in json.loads(result)["error"] + + +@pytest.mark.django_db(transaction=True) +async def test_submit_job_per_repo_prompt_passthrough(): + """Per-repo ``prompt`` on ``RepoSubmitSpec`` reaches the ``RepoTarget`` passed to ``asubmit_batch_runs``.""" + from mcp_server.server import submit_job + from sessions.services import BatchSubmitFailure, BatchSubmitResult + + failure = BatchSubmitFailure(repo_id="a/b", ref="", error="dry-run") + batch_result = BatchSubmitResult(batch_id="00000000-0000-0000-0000-000000000001", runs=[], failed=[failure]) + + submit = AsyncMock(return_value=batch_result) + with patch("mcp_server.server.asubmit_batch_runs", new=submit): + await submit_job(repos=[{"repo_id": "a/b", "prompt": "per-repo override"}], prompt="batch-level") + + targets = submit.call_args.kwargs["repos"] + assert len(targets) == 1 + assert targets[0].prompt == "per-repo override" diff --git a/tests/unit_tests/notifications/test_run_signals.py b/tests/unit_tests/notifications/test_run_signals.py index e17bc746d..01c544716 100644 --- a/tests/unit_tests/notifications/test_run_signals.py +++ b/tests/unit_tests/notifications/test_run_signals.py @@ -222,6 +222,30 @@ def test_user_none_does_not_emit_rollup(self): assert Notification.objects.count() == 0 + def test_delegated_continuation_run_is_silent(self, member_user): + """The coordinator continuation (DELEGATED_JOB, no batch_id) is plumbing — no bell/email.""" + member_user.notify_on_jobs = NotifyOn.ALWAYS + member_user.save(update_fields=["notify_on_jobs"]) + + session = _session(origin=SessionOrigin.MCP_JOB, thread_id=str(uuid.uuid4()), user=member_user) + cont = _run( + session, trigger_type=SessionOrigin.DELEGATED_JOB, user=member_user, continuation_of_batch_id=uuid.uuid4() + ) + run_finished.send(sender=Run, run=cont) + + assert Notification.objects.count() == 0 + + def test_delegated_leg_still_notifies(self, member_user): + """Delegated legs keep their user-facing signal: per-run for a single-leg batch.""" + member_user.notify_on_jobs = NotifyOn.ALWAYS + member_user.save(update_fields=["notify_on_jobs"]) + + session = _session(origin=SessionOrigin.DELEGATED_JOB, thread_id=str(uuid.uuid4()), user=member_user) + leg = _run(session, trigger_type=SessionOrigin.DELEGATED_JOB, user=member_user, batch_id=uuid.uuid4()) + run_finished.send(sender=Run, run=leg) + + assert Notification.objects.filter(recipient=member_user).count() == 1 + def test_webhook_trigger_skipped_before_batch_branch(self, member_user): member_user.notify_on_jobs = NotifyOn.ALWAYS member_user.save(update_fields=["notify_on_jobs"]) diff --git a/tests/unit_tests/sessions/test_continuation.py b/tests/unit_tests/sessions/test_continuation.py new file mode 100644 index 000000000..7447cf1cc --- /dev/null +++ b/tests/unit_tests/sessions/test_continuation.py @@ -0,0 +1,218 @@ +import uuid +from unittest.mock import AsyncMock, patch + +import pytest +from sessions.models import Run, RunStatus, Session, SessionOrigin +from sessions.signals import render_batch_summary, resume_coordinator_on_batch_complete + +pytestmark = pytest.mark.django_db(transaction=True) + + +def _session(thread_id, **kw): + return Session.objects.create( + thread_id=thread_id, origin=SessionOrigin.DELEGATED_JOB, repo_id=kw.pop("repo_id", "g/leg"), **kw + ) + + +def _run(session, **kw): + kw.setdefault("trigger_type", SessionOrigin.DELEGATED_JOB) + kw.setdefault("repo_id", session.repo_id) + kw.setdefault("status", RunStatus.SUCCESSFUL) + return Run.objects.create(session=session, **kw) + + +def test_render_batch_summary_includes_mr_and_truncates(): + batch = uuid.uuid4() + s = _session("leg-1", repo_id="g/a", parent_thread_id="coord") + long_summary = "x" * 600 + runs = [ + _run( + s, + repo_id="g/a", + status=RunStatus.SUCCESSFUL, + batch_id=batch, + merge_request_web_url="https://gl/mr/1", + result_summary=long_summary, + ), + _run( + _session("leg-2", repo_id="g/b", parent_thread_id="coord"), + repo_id="g/b", + status=RunStatus.FAILED, + batch_id=batch, + error_message="boom", + ), + ] + text = render_batch_summary(batch, runs) + assert "1 succeeded, 1 failed" in text + assert "https://gl/mr/1" in text + assert "g/a (successful)" in text + assert "g/b (failed)" in text + # result_summary is truncated to 500 characters by the renderer + assert ("x" * 500) in text + assert ("x" * 501) not in text + + +def test_receiver_ignores_broadcast_batch_without_parent(): + """A batch whose leg session has no parent_thread_id is an ordinary broadcast — no resume.""" + batch = uuid.uuid4() + s = _session("leg-x", parent_thread_id=None) + run = _run(s, batch_id=batch, status=RunStatus.SUCCESSFUL) + resume_coordinator_on_batch_complete(sender=Run, run=run) + assert not Run.objects.filter(continuation_of_batch_id=batch).exists() + + +def test_receiver_creates_one_continuation_when_all_terminal(): + batch = uuid.uuid4() + Session.objects.create(thread_id="coord", origin=SessionOrigin.MCP_JOB, repo_id="g/coord") + leg = _session("leg-1", repo_id="g/a", parent_thread_id="coord") + run = _run(leg, batch_id=batch, status=RunStatus.SUCCESSFUL, merge_request_web_url="https://gl/mr/9") + + with patch("sessions.signals._enqueue_queued_run", return_value=True) as m_enqueue: + resume_coordinator_on_batch_complete(sender=Run, run=run) + + cont = Run.objects.get(continuation_of_batch_id=batch) + assert cont.session_id == "coord" + assert cont.repo_id == "g/coord" + assert cont.trigger_type == SessionOrigin.DELEGATED_JOB + assert "https://gl/mr/9" in cont.prompt + m_enqueue.assert_called_once() + + +def test_receiver_noop_while_a_sibling_is_pending(): + batch = uuid.uuid4() + Session.objects.create(thread_id="coord2", origin=SessionOrigin.MCP_JOB, repo_id="g/coord") + a = _session("leg-a", repo_id="g/a", parent_thread_id="coord2") + b = _session("leg-b", repo_id="g/b", parent_thread_id="coord2") + _run(a, batch_id=batch, status=RunStatus.RUNNING) # still running + done = _run(b, batch_id=batch, status=RunStatus.SUCCESSFUL) + resume_coordinator_on_batch_complete(sender=Run, run=done) + assert not Run.objects.filter(continuation_of_batch_id=batch).exists() + + +def test_receiver_lands_queued_when_coordinator_is_busy(): + """A busy coordinator session (active run) forces the continuation to QUEUED, not READY. + + ``run_one_active_per_session`` rejects a second READY row, so the receiver falls back to + QUEUED for ``dispatch_next_in_session`` to release FIFO — it must not enqueue immediately. + """ + batch = uuid.uuid4() + coord = Session.objects.create(thread_id="coord-busy", origin=SessionOrigin.MCP_JOB, repo_id="g/coord") + # An already-active run on the coordinator session trips the partial unique constraint. + _run(coord, trigger_type=SessionOrigin.MCP_JOB, status=RunStatus.RUNNING) + leg = _session("leg-busy", repo_id="g/a", parent_thread_id="coord-busy") + run = _run(leg, batch_id=batch, status=RunStatus.SUCCESSFUL) + + with patch("sessions.signals._enqueue_queued_run", return_value=True) as m_enqueue: + resume_coordinator_on_batch_complete(sender=Run, run=run) + + cont = Run.objects.get(continuation_of_batch_id=batch) + assert cont.status == RunStatus.QUEUED + assert cont.session_id == "coord-busy" + m_enqueue.assert_not_called() + + +def test_resume_enqueue_failure_is_surfaced(): + """If enqueuing the continuation fails, the coordinator won't auto-resume — log it loudly.""" + batch = uuid.uuid4() + Session.objects.create(thread_id="coord-fail", origin=SessionOrigin.MCP_JOB, repo_id="g/coord") + leg = _session("leg-fail", repo_id="g/a", parent_thread_id="coord-fail") + run = _run(leg, batch_id=batch, status=RunStatus.SUCCESSFUL) + + with patch("sessions.signals.run_job_task") as m_task, patch("sessions.signals.logger") as m_logger: + m_task.aenqueue = AsyncMock(side_effect=RuntimeError("broker down")) + resume_coordinator_on_batch_complete(sender=Run, run=run) + + # The continuation row exists (FAILED, awaiting the retry sweep) and the failure is not silent. + cont = Run.objects.get(continuation_of_batch_id=batch) + assert cont.status == RunStatus.FAILED + assert m_logger.error.called + assert "auto-resume" in m_logger.error.call_args[0][0].lower() + + +def test_continuation_is_created_queued_before_promotion(): + """QUEUED is the crash-safe resting state: a crash after the insert leaves a row the + orphan-release cron can recover, never a stuck READY row nothing owns.""" + batch = uuid.uuid4() + Session.objects.create(thread_id="coord-rest", origin=SessionOrigin.MCP_JOB, repo_id="g/coord") + leg = _session("leg-rest", repo_id="g/a", parent_thread_id="coord-rest") + run = _run(leg, batch_id=batch, status=RunStatus.SUCCESSFUL) + + with patch("sessions.signals.release_next_queued") as m_release: + resume_coordinator_on_batch_complete(sender=Run, run=run) + + cont = Run.objects.get(continuation_of_batch_id=batch) + assert cont.status == RunStatus.QUEUED + m_release.assert_called_once_with("coord-rest") + + +def test_receiver_leaves_continuation_queued_behind_active_chat_run(): + """An active chat run is invisible to run_one_active_per_session; the any-active guard + must still keep the continuation QUEUED instead of enqueuing it into the held session lock.""" + batch = uuid.uuid4() + coord = Session.objects.create(thread_id="coord-chat", origin=SessionOrigin.CHAT, repo_id="g/coord") + _run(coord, trigger_type=SessionOrigin.CHAT, status=RunStatus.RUNNING) + leg = _session("leg-chat", repo_id="g/a", parent_thread_id="coord-chat") + run = _run(leg, batch_id=batch, status=RunStatus.SUCCESSFUL) + + with patch("sessions.signals._enqueue_queued_run", return_value=True) as m_enqueue: + resume_coordinator_on_batch_complete(sender=Run, run=run) + + cont = Run.objects.get(continuation_of_batch_id=batch) + assert cont.status == RunStatus.QUEUED + m_enqueue.assert_not_called() + + +def test_receiver_ignores_non_delegated_batch_runs(): + """API/MCP broadcast batches never resume anything — the trigger-type guard fires first.""" + batch = uuid.uuid4() + s = Session.objects.create(thread_id="api-1", origin=SessionOrigin.API_JOB, repo_id="g/x") + run = Run.objects.create( + session=s, repo_id="g/x", trigger_type=SessionOrigin.API_JOB, status=RunStatus.SUCCESSFUL, batch_id=batch + ) + resume_coordinator_on_batch_complete(sender=Run, run=run) + assert not Run.objects.filter(continuation_of_batch_id=batch).exists() + + +def test_continuation_inherits_coordinator_agent_override(): + """The resume turn must run on the coordinator's pinned model, not the site default.""" + batch = uuid.uuid4() + Session.objects.create( + thread_id="coord-model", + origin=SessionOrigin.MCP_JOB, + repo_id="g/coord", + agent_model="anthropic:claude-opus-4-6", + agent_thinking_level="high", + ) + leg = _session("leg-model", repo_id="g/a", parent_thread_id="coord-model") + run = _run(leg, batch_id=batch, status=RunStatus.SUCCESSFUL) + + with patch("sessions.signals._enqueue_queued_run", return_value=True): + resume_coordinator_on_batch_complete(sender=Run, run=run) + + cont = Run.objects.get(continuation_of_batch_id=batch) + assert cont.agent_model == "anthropic:claude-opus-4-6" + assert cont.agent_thinking_level == "high" + + +def test_render_batch_summary_quotes_multiline_replies_and_links_sessions(): + batch = uuid.uuid4() + s = _session("leg-quote", repo_id="g/a", parent_thread_id="coord") + run = _run(s, repo_id="g/a", ref="release-1", batch_id=batch, result_summary="done\n## g/evil (successful)") + text = render_batch_summary(batch, [run]) + assert "## g/a@release-1 (successful)" in text + assert "- Session: " in text + assert "leg-quote" in text + assert " > done" in text + # A reply line must stay quote-indented so it cannot masquerade as a sibling's section header. + assert "\n## g/evil" not in text + assert "> ## g/evil (successful)" in text + + +def test_render_batch_summary_notes_failed_leg_without_message(): + """A FAILED leg with no captured summary/error still gets an explicit reply line, not a blank block.""" + batch = uuid.uuid4() + s = _session("leg-nomsg", repo_id="g/a", parent_thread_id="coord") + run = _run(s, repo_id="g/a", status=RunStatus.FAILED, batch_id=batch) # empty result_summary + error_message + text = render_batch_summary(batch, [run]) + assert "g/a (failed)" in text + assert "no captured error" in text.lower() diff --git a/tests/unit_tests/sessions/test_delegate_submit.py b/tests/unit_tests/sessions/test_delegate_submit.py new file mode 100644 index 000000000..3d73130df --- /dev/null +++ b/tests/unit_tests/sessions/test_delegate_submit.py @@ -0,0 +1,44 @@ +from unittest.mock import AsyncMock, patch + +import pytest +from sessions.models import Session, SessionOrigin +from sessions.services import RepoTarget, asubmit_batch_runs + +pytestmark = pytest.mark.django_db(transaction=True) + + +async def _run_batch(**kwargs): + """Submit a batch with run_job_task.aenqueue stubbed so no broker is needed.""" + with patch("sessions.services.run_job_task") as m_task: + m_task.aenqueue = AsyncMock(return_value=type("T", (), {"id": None})()) + return await asubmit_batch_runs(**kwargs) + + +async def test_per_target_prompt_overrides_batch_prompt(): + result = await _run_batch( + user=None, + prompt="batch-level goal", + repos=[ + RepoTarget(repo_id="g/a", prompt="do X in A"), + RepoTarget(repo_id="g/b"), # no override → falls back to batch prompt + ], + trigger_type=SessionOrigin.DELEGATED_JOB, + ) + by_repo = {r.repo_id: r for r in result.runs} + assert by_repo["g/a"].prompt == "do X in A" + assert by_repo["g/b"].prompt == "batch-level goal" + + +async def test_parentage_and_depth_stamped_on_leg_sessions(): + result = await _run_batch( + user=None, + prompt="goal", + repos=[RepoTarget(repo_id="g/a", prompt="p")], + trigger_type=SessionOrigin.DELEGATED_JOB, + parent_thread_id="parent-thread-123", + spawn_depth=2, + ) + leg = result.runs[0] + session = await Session.objects.aget(thread_id=leg.session_id) + assert session.parent_thread_id == "parent-thread-123" + assert session.spawn_depth == 2 diff --git a/tests/unit_tests/sessions/test_management.py b/tests/unit_tests/sessions/test_management.py index 2dd7623e9..a4c75c2ab 100644 --- a/tests/unit_tests/sessions/test_management.py +++ b/tests/unit_tests/sessions/test_management.py @@ -195,6 +195,78 @@ def test_releases_queued_when_no_active_sibling(self, create_db_task_result): assert orphan.task_result_id == fake_task.id assert "Released: 1" in out.getvalue() + def test_requeues_and_releases_never_started_failed_continuation(self, create_db_task_result): + """A continuation that FAILED before ever starting (dispatch failure, no linked task) + is the batch's only shot at resuming its coordinator — re-queue and release it.""" + coord = _make_session() + cont = Run.objects.create( + session=coord, + trigger_type=SessionOrigin.DELEGATED_JOB, + repo_id="a/b", + status=RunStatus.FAILED, + continuation_of_batch_id=uuid.uuid4(), + error_message="dispatch_failed: RuntimeError: broker down", + prompt="p", + ) + fake_task = MagicMock(id=create_db_task_result().id) + out = StringIO() + with patch("sessions.signals.run_job_task") as mock_task: + mock_task.aenqueue = AsyncMock(return_value=fake_task) + call_command("release_orphan_queued_sessions", stdout=out) + + cont.refresh_from_db() + assert cont.status == RunStatus.READY + assert cont.error_message == "" + assert cont.task_result_id == fake_task.id + assert "requeued continuations: 1" in out.getvalue() + + def test_does_not_requeue_continuation_that_ran(self, create_db_task_result): + """A continuation that actually executed and FAILED must not loop forever.""" + coord = _make_session() + cont = Run.objects.create( + session=coord, + trigger_type=SessionOrigin.DELEGATED_JOB, + repo_id="a/b", + status=RunStatus.FAILED, + continuation_of_batch_id=uuid.uuid4(), + task_result_id=create_db_task_result().id, + error_message="dispatch_failed: TimeoutError: session lock", + prompt="p", + ) + out = StringIO() + with patch("sessions.signals.run_job_task") as mock_task: + mock_task.aenqueue = AsyncMock() + call_command("release_orphan_queued_sessions", stdout=out) + mock_task.aenqueue.assert_not_called() + + cont.refresh_from_db() + assert cont.status == RunStatus.FAILED + assert "requeued continuations: 0" in out.getvalue() + + def test_does_not_requeue_link_failed_continuation(self): + """link_failed means the agent task exists and will run — re-queuing would double-run it.""" + from sessions.signals import LINK_FAILED_PREFIX + + coord = _make_session() + cont = Run.objects.create( + session=coord, + trigger_type=SessionOrigin.DELEGATED_JOB, + repo_id="a/b", + status=RunStatus.FAILED, + continuation_of_batch_id=uuid.uuid4(), + error_message=f"{LINK_FAILED_PREFIX}: RuntimeError: db blip", + prompt="p", + ) + out = StringIO() + with patch("sessions.signals.run_job_task") as mock_task: + mock_task.aenqueue = AsyncMock() + call_command("release_orphan_queued_sessions", stdout=out) + mock_task.aenqueue.assert_not_called() + + cont.refresh_from_db() + assert cont.status == RunStatus.FAILED + assert "requeued continuations: 0" in out.getvalue() + def test_skips_queued_when_active_sibling_exists(self): """A QUEUED run whose session already has a READY/RUNNING sibling is left alone.""" session = _make_session() diff --git a/tests/unit_tests/sessions/test_models.py b/tests/unit_tests/sessions/test_models.py index 438750da2..6acb1c186 100644 --- a/tests/unit_tests/sessions/test_models.py +++ b/tests/unit_tests/sessions/test_models.py @@ -35,6 +35,7 @@ def test_session_origin_includes_chat(): "ui_job", "issue_webhook", "mr_webhook", + "delegated_job", } @@ -131,7 +132,11 @@ def test_active_constraint_literals_match_enums(): constraint = next(c for c in Run._meta.constraints if c.name == "run_one_active_per_session") conditions = dict(constraint.condition.children) assert set(conditions["status__in"]) == {RunStatus.READY, RunStatus.RUNNING} - assert set(conditions["trigger_type__in"]) == {SessionOrigin.API_JOB, SessionOrigin.MCP_JOB} + assert set(conditions["trigger_type__in"]) == { + SessionOrigin.API_JOB, + SessionOrigin.MCP_JOB, + SessionOrigin.DELEGATED_JOB, + } def test_session_origin_check_constraint_rejects_unknown_value(): @@ -251,6 +256,41 @@ def test_effective_notify_on_falls_back_to_never_without_override_schedule_or_us assert run.effective_notify_on == NotifyOn.NEVER +# --- delegate_jobs data model ------------------------------------------- + + +def test_delegated_job_is_accepted_by_enum_check_constraints(): + session = _mk_session(origin=SessionOrigin.DELEGATED_JOB) + _mk_run(session, trigger_type=SessionOrigin.DELEGATED_JOB, status=RunStatus.QUEUED) + + +def test_one_continuation_per_batch_enforced(): + batch = uuid.uuid4() + session = _mk_session() + _mk_run(session, status=RunStatus.QUEUED, continuation_of_batch_id=batch) + with pytest.raises(IntegrityError): + _mk_run(session, status=RunStatus.QUEUED, continuation_of_batch_id=batch) + + +def test_null_continuation_of_batch_not_deduplicated(): + """Multiple runs with NULL continuation_of_batch_id coexist (partial constraint).""" + session = _mk_session() + _mk_run(session, status=RunStatus.QUEUED) + _mk_run(session, status=RunStatus.QUEUED) # no IntegrityError + + +def test_spawn_depth_cannot_exceed_cap(): + """The delegation-depth fuse is enforced at the DB, not only in the delegate_jobs tool.""" + from sessions.models import MAX_SPAWN_DEPTH + + _mk_session(spawn_depth=MAX_SPAWN_DEPTH) # at the cap is allowed + with pytest.raises(IntegrityError): + _mk_session(spawn_depth=MAX_SPAWN_DEPTH + 1) + + +# --- message_id ------------------------------------------------------------ + + def test_run_message_id_defaults_blank_and_persists(session_fixture): run = Run.objects.create( session=session_fixture, diff --git a/tests/unit_tests/sessions/test_services.py b/tests/unit_tests/sessions/test_services.py index 072f8a0a7..5ba11e03f 100644 --- a/tests/unit_tests/sessions/test_services.py +++ b/tests/unit_tests/sessions/test_services.py @@ -152,6 +152,26 @@ async def test_submit_continuation_queues_when_thread_busy(): assert await Run.objects.filter(session_id=tid).acount() == 2 +@pytest.mark.django_db(transaction=True) +async def test_submit_batch_creates_all_rows_before_first_enqueue(): + """A leg can turn terminal (emitting run_finished) as soon as it is enqueued, so every + sibling row must already exist by then — else batch-completion receivers elect on a subset.""" + seen_batch_counts: list[int] = [] + + async def _fake_enqueue(**kwargs): + run = await Run.objects.aget(pk=kwargs["run_id"]) + seen_batch_counts.append(await Run.objects.filter(batch_id=run.batch_id).acount()) + return await _atask_result_row(uuid.uuid4()) + + repos = [RepoTarget(repo_id=f"g/r{i}") for i in range(3)] + with patch("sessions.services.run_job_task") as mock_task: + mock_task.aenqueue = AsyncMock(side_effect=_fake_enqueue) + result = await asubmit_batch_runs(user=None, prompt="p", repos=repos, trigger_type=SessionOrigin.API_JOB) + + assert len(result.runs) == 3 + assert seen_batch_counts == [3, 3, 3] + + # --------------------------------------------------------------------------- # validate_repo_list tests (ported from activity) # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/sessions/test_tasks.py b/tests/unit_tests/sessions/test_tasks.py index 8645b4464..a6b0d3398 100644 --- a/tests/unit_tests/sessions/test_tasks.py +++ b/tests/unit_tests/sessions/test_tasks.py @@ -21,6 +21,16 @@ def test_sync_stuck_runs_cron_task_dispatches_command(): mock_call_command.assert_called_once_with("sync_stuck_runs") +def test_release_orphan_queued_sessions_cron_task_dispatches_command(): + """The cron task dispatches the release_orphan_queued_sessions management command.""" + from sessions.tasks import release_orphan_queued_sessions_cron_task + + with patch("sessions.tasks.call_command") as mock_call_command: + release_orphan_queued_sessions_cron_task.func() + + mock_call_command.assert_called_once_with("release_orphan_queued_sessions") + + # --- reclassify_missing_envelopes_cron_task (Epic 1 review backstop) --------