secondshift 0.1.0b12__py3-none-any.whl

secondshift/runs/deliveries.py

@@ -0,0 +1,80 @@
+"""Delivering a finished run's results again, without running the agent again."""
+
+import uuid
+
+from dbos import EnqueueOptions
+from django.conf import settings
+from django.db import transaction
+
+from secondshift.config.constants import (
+    RETRY_ISSUES_WORKFLOW,
+    RETRY_NOTIFICATION_WORKFLOW,
+    RETRY_PULL_REQUEST_WORKFLOW,
+)
+from secondshift.runs.dao import add_event, lock_delivery
+from secondshift.runs.models import Run, RunDelivery
+from secondshift.runs.orchestration.dbos_runtime import new_dbos_client
+from secondshift.runs.services import RunRefused
+
+
+class DeliveryQueueUnavailable(Exception):
+    pass
+
+
+def retryable(delivery: RunDelivery) -> bool:
+    # Pull requests and issues need no connection. The repository's own
+    # GitHub access sends them.
+    return delivery.status == RunDelivery.Status.FAILED and (
+        delivery.kind in {RunDelivery.Kind.PULL_REQUEST, RunDelivery.Kind.ISSUES}
+        or delivery.connection_id is not None
+    )
+
+
+@transaction.atomic
+def retry_delivery(run: Run, delivery_id: int) -> RunDelivery:
+    """Queue a failed delivery again. Authorization is checked when it is sent."""
+    delivery = lock_delivery(run, delivery_id)
+    if delivery.status != RunDelivery.Status.FAILED:
+        raise RunRefused("Only a failed delivery can be retried.")
+    if not retryable(delivery):
+        raise RunRefused("The connection for this delivery was removed.")
+    delivery.status = RunDelivery.Status.RETRYING
+    delivery.message = "Waiting to retry."
+    delivery.save(update_fields=["status", "message", "updated_at"])
+    # Enqueued inside the transaction, so a queue outage leaves it failed.
+    if delivery.kind == RunDelivery.Kind.PULL_REQUEST:
+        add_event(run.pk, "Retrying the pull request.")
+        enqueue_retry(RETRY_PULL_REQUEST_WORKFLOW, f"pull-request-{run.pk}", run.pk)
+    elif delivery.kind == RunDelivery.Kind.ISSUES:
+        add_event(run.pk, "Retrying the issues.")
+        enqueue_retry(RETRY_ISSUES_WORKFLOW, f"issues-{run.pk}", run.pk)
+    else:
+        add_event(run.pk, f"Retrying the notification to {delivery.target}.")
+        enqueue_retry(
+            RETRY_NOTIFICATION_WORKFLOW,
+            f"{run.pk}-{delivery.connection_id}",
+            run.pk,
+            delivery.connection_id,
+        )
+    return delivery
+
+
+def enqueue_retry(workflow: str, name: str, *args: int) -> str:
+    suffix = uuid.uuid4().hex[:12]
+    options: EnqueueOptions = {
+        "workflow_name": workflow,
+        "queue_name": settings.DBOS_QUEUE_NAME,
+        # Each retry is its own workflow. A reused ID would return the old result.
+        "workflow_id": f"secondshift-delivery-{name}-{suffix}",
+        "application_name": settings.DBOS_APPLICATION_NAME,
+    }
+    try:
+        client = new_dbos_client()
+    except Exception as error:
+        raise DeliveryQueueUnavailable from error
+    try:
+        return client.enqueue(options, *args).get_workflow_id()
+    except Exception as error:
+        raise DeliveryQueueUnavailable from error
+    finally:
+        client.destroy()

secondshift/runs/execution/access_needs.py

@@ -0,0 +1,89 @@
+"""Let the agent say what access it needed and this run did not have.
+
+A skill from another repository often expects services nobody connected. The
+agent cannot see tools it was not given, so it reports them here instead of
+quietly doing less. The run page lists them. Reporting never grants anything.
+"""
+
+from typing import Annotated, Any
+
+from pydantic import Field
+from pydantic_ai import ModelRetry, RunContext, Tool
+
+from secondshift.runs.dao import (
+    access_need_for_call,
+    create_access_need,
+    run_access_needs,
+    update_access_need_tools,
+)
+from secondshift.runs.execution.workspace_tools import WorkspaceToolDeps
+
+TOOL_NAME = "report_missing_access"
+# Enough for a skill that names several services, without filling the page.
+MAX_SERVICES = 20
+MAX_TOOLS = 30
+
+ServiceArgument = Annotated[str, Field(min_length=1, max_length=100)]
+ToolsArgument = Annotated[
+    list[Annotated[str, Field(min_length=1, max_length=100)]],
+    Field(max_length=MAX_TOOLS),
+]
+ReasonArgument = Annotated[str, Field(max_length=1000)]
+
+ACCESS_NOTE = f"""If the task needs a service, tool or credential this run does not have,
+such as an MCP server or API the skill names, call {TOOL_NAME} once for each
+service, then go on with what you can do. Do not try to reach the service
+another way. Name the missing access in your report too."""
+
+
+def report_missing_access(
+    ctx: RunContext[WorkspaceToolDeps],
+    service: ServiceArgument,
+    tools: ToolsArgument | None = None,
+    reason: ReasonArgument = "",
+) -> dict[str, Any]:
+    """Record a service or tool this task needs and this run was not given.
+
+    Name the service, such as linear or sentry, the tools needed from it, such
+    as create_issue, and why. People grant access after the run. Nothing is
+    granted now.
+    """
+    if not ctx.tool_call_id:
+        raise RuntimeError("The model tool call has no stable identifier.")
+    service = " ".join(service.split())
+    if not service:
+        raise ModelRetry("Name the service.")
+    names = sorted({" ".join(tool.split()) for tool in tools or []} - {""})
+    run_id = ctx.deps.run_id
+    if access_need_for_call(run_id, ctx.tool_call_id):
+        return recorded()
+    needs = run_access_needs(run_id)
+    same = next(
+        (need for need in needs if need.service.casefold() == service.casefold()),
+        None,
+    )
+    if same is not None:
+        merged = sorted(set(same.tools) | set(names))[:MAX_TOOLS]
+        if merged != same.tools:
+            update_access_need_tools(same, merged)
+        return recorded()
+    if len(needs) >= MAX_SERVICES:
+        return {
+            "recorded": False,
+            "note": f"This run already reported {MAX_SERVICES} services. "
+            "Name the rest in the report.",
+        }
+    create_access_need(run_id, ctx.tool_call_id, service, names, reason.strip())
+    return recorded()
+
+
+def recorded() -> dict[str, Any]:
+    return {
+        "recorded": True,
+        "note": "Recorded for people to review after the run. It is not "
+        "available in this run. Continue without it.",
+    }
+
+
+def access_tool() -> Tool[WorkspaceToolDeps]:
+    return Tool(report_missing_access, name=TOOL_NAME)

secondshift/runs/execution/agent.py

@@ -0,0 +1,91 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from asgiref.sync import sync_to_async
+from pydantic_ai import Agent, UsageLimits
+from pydantic_ai.capabilities import ResolveModelId
+from pydantic_ai.durable_exec.dbos import DBOSDurability
+from pydantic_ai.models import Model, ModelResolutionContext
+
+from secondshift.runs.dao import get_run
+from secondshift.runs.execution.budgets import FinishWithinBudget
+from secondshift.runs.execution.deadlines import check_deadline
+from secondshift.runs.execution.model_gateway import (
+    issue_model_gateway_access,
+    pydantic_model,
+)
+
+RUN_MODEL_ID = "secondshift-run-model"
+
+
+@dataclass(frozen=True)
+class AgentLimits:
+    max_turns: int
+    max_tool_calls: int
+    max_tool_result_chars: int = 65_536
+
+    def __post_init__(self) -> None:
+        if self.max_turns < 1:
+            raise ValueError("max_turns must be positive")
+        if self.max_tool_calls < 0:
+            raise ValueError("max_tool_calls cannot be negative")
+        if self.max_tool_result_chars < 1:
+            raise ValueError("max_tool_result_chars must be positive")
+
+    def usage_limits(self) -> UsageLimits:
+        """Translate server policy to Pydantic AI's native run limits."""
+        return UsageLimits(
+            request_limit=self.max_turns,
+            tool_calls_limit=self.max_tool_calls,
+        )
+
+
+async def resolve_run_model(
+    context: ModelResolutionContext[Any], model_id: str
+) -> Model | None:
+    """Build gateway access inside the durable model-request step.
+
+    Only the run ID crosses the DBOS boundary. The short-lived gateway token is
+    created and consumed inside the step and is never serialized into workflow
+    inputs, outputs, or agent dependencies.
+    """
+    if model_id != RUN_MODEL_ID:
+        return None
+    check_deadline(context.deps.deadline)
+    # Pydantic AI resolves models inside its event loop, where Django refuses
+    # synchronous database access.
+    return await sync_to_async(run_model)(int(context.deps.run_id))
+
+
+def run_model(run_id: int) -> Model:
+    run = get_run(run_id, "job")
+    return pydantic_model(issue_model_gateway_access(run))
+
+
+def build_agent(
+    *,
+    name: str = "secondshift",
+    instructions: str,
+    tools: Sequence[Any] = (),
+    deps_type: type[Any] = object,
+) -> Agent[Any, str]:
+    """Build the durable orchestrator-side agent for one run.
+
+    Tools supplied here are server-built wrappers around the runner or capability
+    broker. Repository content cannot register tools or change durability.
+    """
+    return Agent(
+        model=RUN_MODEL_ID,
+        name=name,
+        instructions=instructions,
+        deps_type=deps_type,
+        tools=tools,
+        defer_model_check=True,
+        capabilities=[
+            ResolveModelId(resolve_run_model),
+            FinishWithinBudget(),
+            # Durable steps are named by the agent, so each agent needs its own.
+            DBOSDurability(name=name),
+        ],
+    )

secondshift/runs/execution/agent_broker.py

@@ -0,0 +1,198 @@
+"""Stateless MCP transport over the existing, run-scoped tool policy."""
+
+import hashlib
+import json
+from typing import Any
+
+from django.db import transaction
+from django.http import HttpRequest, HttpResponse, JsonResponse
+from django.views.decorators.csrf import csrf_exempt
+from pydantic_ai import ModelRetry, RunContext, RunUsage, Tool, ToolFailed
+
+from secondshift.connections.agents import usable_agent_connection
+from secondshift.connections.dao import find_agent_connection
+from secondshift.connections.models import AgentConnection
+from secondshift.runs.artifacts import run_payload
+from secondshift.runs.dao import active_lease, get_run, lock_lease
+from secondshift.runs.execution.access_needs import access_tool
+from secondshift.runs.execution.deadlines import check_deadline, run_deadline
+from secondshift.runs.execution.external_tools import external_tools
+from secondshift.runs.execution.issues import issue_limit, issue_tool
+from secondshift.runs.execution.payload import payload_tool
+from secondshift.runs.execution.workspace_tools import (
+    MODEL_TOOL_NAMES,
+    WorkspaceToolDeps,
+    workspace_tools,
+)
+from secondshift.runs.models import AgentLease, Run
+
+# Streamable HTTP MCP versions the broker speaks, newest first.
+PROTOCOL_VERSIONS = ("2025-11-25", "2025-06-18", "2025-03-26")
+
+
+class AgentAccessError(ValueError):
+    pass
+
+
+class MethodNotFound(ValueError):
+    """Clients probe newer methods first and fall back when they are missing."""
+
+
+def run_connection(run: Run) -> AgentConnection:
+    plan = run.plan.execution
+    connection = find_agent_connection(
+        plan.connection_id,
+        run.workspace_id,
+        plan.owner_id,
+        plan.credential_key,
+        plan.backend,
+    )
+    if connection is None or not usable_agent_connection(connection, run.workspace_id):
+        raise AgentAccessError(
+            "The agent login is no longer available. Reconnect and start a new run."
+        )
+    return connection
+
+
+def lease_run(lease: AgentLease) -> Run:
+    run = get_run(lease.run_id, "job")
+    if not lease.active or run.status != Run.Status.RUNNING:
+        raise AgentAccessError("This agent run is no longer active.")
+    check_deadline(run_deadline(run))
+    run_connection(run)
+    return run
+
+
+def broker_tools(run: Run) -> dict[str, Tool[WorkspaceToolDeps]]:
+    allowed = set(run.plan.agent.tools)
+    local = {MODEL_TOOL_NAMES[name] for name in allowed if name in MODEL_TOOL_NAMES}
+    tools = [tool for tool in workspace_tools() if tool.name in local]
+    if run_payload(run):
+        tools.append(payload_tool())
+    if issue_limit(run.plan):
+        tools.append(issue_tool())
+    tools.append(access_tool())
+    return {tool.name: tool for tool in (*tools, *external_tools(run))}
+
+
+def tool_result(
+    lease: AgentLease, name: str, arguments: dict[str, Any]
+) -> dict[str, Any]:
+    with transaction.atomic():
+        lease = lock_lease(lease.pk)
+        run = lease_run(lease)
+        if lease.tool_calls >= run.plan.agent.limits.max_tool_calls:
+            raise AgentAccessError("This run has used its tool call budget.")
+        # Resolve and validate before counting the call.
+        tool = broker_tools(run).get(name)
+        if tool is None:
+            raise AgentAccessError("This tool is not approved for the run.")
+        validated = tool.function_schema.validator.validate_python(arguments)
+        lease.tool_calls += 1
+        lease.save(update_fields=["tool_calls"])
+    # The broker numbers calls itself. Clients restart their request IDs when
+    # they reconnect, so those IDs cannot tell two calls apart.
+    call_id = f"agent-{lease.pk}-{lease.tool_calls}"
+    deps = WorkspaceToolDeps(
+        run_id=run.pk,
+        session_id=lease.session_id,
+        max_output_chars=run.plan.agent.limits.max_tool_result_chars,
+        deadline=run_deadline(run),
+        allowed_tools=run.plan.agent.tools,
+        payload=run_payload(run),
+        issue_limit=issue_limit(run.plan),
+    )
+    context = RunContext(deps=deps, model=None, usage=RunUsage(), tool_call_id=call_id)
+    try:
+        value = tool.function(context, **validated)
+        return {
+            "content": [{"type": "text", "text": json.dumps(value)}],
+            "isError": False,
+        }
+    except (ModelRetry, ToolFailed) as error:
+        # SecondShift wrote these for the agent, for example which argument to fix.
+        return {"content": [{"type": "text", "text": error.message}], "isError": True}
+    except Exception:
+        # Provider exceptions can include request contents. Return no raw errors.
+        return {
+            "content": [
+                {"type": "text", "text": "The approved tool could not complete."}
+            ],
+            "isError": True,
+        }
+
+
+def mcp_request(lease: AgentLease, message: dict[str, Any]) -> dict[str, Any] | None:
+    run = lease_run(lease)
+    method = message.get("method")
+    if "id" not in message:
+        return None
+    if method == "initialize":
+        requested = (message.get("params") or {}).get("protocolVersion")
+        return {
+            "protocolVersion": requested
+            if requested in PROTOCOL_VERSIONS
+            else PROTOCOL_VERSIONS[0],
+            "capabilities": {"tools": {}},
+            "serverInfo": {"name": "secondshift-run", "version": "1.0"},
+        }
+    if method == "ping":
+        return {}
+    if method == "tools/list":
+        return {
+            "tools": [
+                {
+                    "name": tool.name,
+                    "description": tool.description or tool.name,
+                    "inputSchema": tool.function_schema.json_schema,
+                }
+                for tool in broker_tools(run).values()
+            ]
+        }
+    if method == "tools/call":
+        params = message.get("params", {})
+        arguments = params.get("arguments", {})
+        if not isinstance(arguments, dict):
+            raise AgentAccessError("Invalid tool arguments.")
+        return tool_result(lease, str(params.get("name", "")), arguments)
+    raise MethodNotFound(str(method))
+
+
+@csrf_exempt
+def agent_mcp(request: HttpRequest) -> HttpResponse:
+    if request.method != "POST":
+        return HttpResponse(status=405)
+    token = request.headers.get("Authorization", "").removeprefix("Bearer ")
+    lease = active_lease(hashlib.sha256(token.encode()).hexdigest()) if token else None
+    if lease is None:
+        return HttpResponse(status=401)
+    if len(request.body) > 1024 * 1024:
+        return HttpResponse(status=413)
+    message: dict[str, Any] = {}
+    try:
+        message = json.loads(request.body)
+        if not isinstance(message, dict):
+            return HttpResponse(status=400)
+        result = mcp_request(lease, message)
+        if "id" not in message:
+            return HttpResponse(status=202)
+        return JsonResponse({"jsonrpc": "2.0", "id": message["id"], "result": result})
+    except MethodNotFound:
+        return JsonResponse(
+            {
+                "jsonrpc": "2.0",
+                "id": message.get("id"),
+                "error": {"code": -32601, "message": "Method not found."},
+            }
+        )
+    except Exception:
+        return JsonResponse(
+            {
+                "jsonrpc": "2.0",
+                "id": message.get("id"),
+                "error": {
+                    "code": -32000,
+                    "message": "Agent access denied or tool request invalid.",
+                },
+            }
+        )

secondshift/runs/execution/budgets.py

@@ -0,0 +1,177 @@
+"""Run budgets: model requests, tool calls and tokens.
+
+The agent sees when it has to finish. On its last allowed model request, or
+once its tool calls are used up, its tools are hidden and it is told to write
+the report with what it has. A run that still runs out keeps what the model
+wrote so far as a partial report. So does a run that times out.
+"""
+
+import datetime
+from collections import Counter
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import Any
+
+from pydantic_ai import RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.exceptions import ModelHTTPError
+from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.usage import RunUsage
+
+from secondshift.config.constants import TOKEN_BUDGET_MESSAGE, AgentBackend
+from secondshift.jobs.plan import ExecutionPlan
+from secondshift.runs.models import RunFailure
+from secondshift.runs.reporting.failures import StageFailure
+from secondshift.workspaces.limits import RunLimit, raise_hint
+
+BUDGET_HEADLINE = RunFailure.BUDGET.label
+FINISH_INSTRUCTIONS = (
+    "This run has reached its budget, so no tools are available. Write the final "
+    "report now from what you have found. Say what you could not check."
+)
+PARTIAL_REPORT_NOTE = (
+    "_The run stopped at one of its limits before the agent finished. "
+    "This is what it wrote so far._"
+)
+
+
+class Budget(StrEnum):
+    REQUESTS = "requests"
+    TOOL_CALLS = "tool_calls"
+    TOKENS = "tokens"
+
+
+@dataclass(frozen=True)
+class BudgetLimit:
+    label: str
+    limit: RunLimit
+
+
+BUDGET_LIMITS = {
+    Budget.REQUESTS: BudgetLimit("model requests", RunLimit.MAX_TURNS),
+    Budget.TOOL_CALLS: BudgetLimit("tool calls", RunLimit.MAX_TOOL_CALLS),
+    Budget.TOKENS: BudgetLimit("tokens", RunLimit.TOKEN_BUDGET),
+}
+
+
+def spent_budget(context: RunContext[Any]) -> Budget | None:
+    """The budget that leaves the agent no room for another tool round."""
+    limits = context.usage_limits
+    if limits is None:
+        return None
+    usage = context.usage
+    # The request about to be made is not counted yet.
+    if limits.request_limit is not None and usage.requests + 1 >= limits.request_limit:
+        return Budget.REQUESTS
+    if (
+        limits.tool_calls_limit is not None
+        and usage.tool_calls >= limits.tool_calls_limit
+    ):
+        return Budget.TOOL_CALLS
+    return None
+
+
+class FinishWithinBudget(AbstractCapability[Any]):
+    """Hide the tools and ask for the report once the budget is spent."""
+
+    async def prepare_tools(
+        self, ctx: RunContext[Any], tool_defs: list[ToolDefinition]
+    ) -> list[ToolDefinition]:
+        return [] if spent_budget(ctx) else tool_defs
+
+    def get_instructions(self) -> Any:
+        def instructions(ctx: RunContext[Any]) -> str:
+            return FINISH_INSTRUCTIONS if spent_budget(ctx) else ""
+
+        return instructions
+
+
+def reached_budget(
+    usage: RunUsage, max_turns: int, max_tool_calls: int
+) -> Budget | None:
+    """The budget a finished run used in full, if any."""
+    if usage.requests >= max_turns:
+        return Budget.REQUESTS
+    if usage.tool_calls >= max_tool_calls:
+        return Budget.TOOL_CALLS
+    return None
+
+
+def exhausted_budget(messages: Sequence[ModelMessage], max_turns: int) -> Budget:
+    """Which native limit stopped the run, from the messages it produced."""
+    responses = sum(isinstance(message, ModelResponse) for message in messages)
+    return Budget.REQUESTS if responses >= max_turns else Budget.TOOL_CALLS
+
+
+def is_token_budget_refusal(error: BaseException) -> bool:
+    if not isinstance(error, ModelHTTPError) or error.status_code != 429:
+        return False
+    return TOKEN_BUDGET_MESSAGE in str(error.body)
+
+
+def budget_reason(budget: Budget, limit: int) -> str:
+    spec = BUDGET_LIMITS[budget]
+    return (
+        f"Stopped at {limit:,} {spec.label}, the run's limit. "
+        f"{raise_hint(spec.limit, limit)}"
+    )
+
+
+def budget_reached_message(budget: Budget) -> str:
+    return (
+        f"The run used all its {BUDGET_LIMITS[budget].label}, so the agent "
+        "wrote its report without further tool calls."
+    )
+
+
+def budget_failure(budget: Budget, limit: int) -> StageFailure:
+    reason = budget_reason(budget, limit)
+    return StageFailure(RunFailure.BUDGET, reason, f"{BUDGET_HEADLINE}: {reason}")
+
+
+def partial_report(messages: Sequence[ModelMessage]) -> str:
+    """What the model wrote and which tools it called before it was stopped."""
+    texts: list[str] = []
+    calls: Counter[str] = Counter()
+    for message in messages:
+        if not isinstance(message, ModelResponse):
+            continue
+        for part in message.parts:
+            if isinstance(part, TextPart) and part.content.strip():
+                texts.append(part.content.strip())
+            elif isinstance(part, ToolCallPart):
+                calls[part.tool_name] += 1
+    if not texts and not calls:
+        return ""
+    sections = [PARTIAL_REPORT_NOTE, *texts]
+    if calls:
+        sections.append(
+            "Tool calls made:\n"
+            + "\n".join(f"- `{name}`: {count}" for name, count in calls.items())
+        )
+    return "\n\n".join(sections)
+
+
+def limits_note(
+    plan: ExecutionPlan, backend: AgentBackend, started_at: datetime.datetime | None
+) -> str:
+    """The prompt lines that tell the agent how much room the run has."""
+    minutes = plan.limit(RunLimit.TIMEOUT_SECONDS) // 60
+    limits = [
+        f"{minutes} minutes",
+        f"{plan.limit(RunLimit.MAX_TOOL_CALLS):,} tool calls",
+    ]
+    if backend.limits_model_calls:
+        limits.append(f"{plan.limit(RunLimit.MAX_TURNS):,} model calls")
+    start = ""
+    if started_at is not None:
+        start = started_at.astimezone(datetime.UTC).strftime(
+            " It started at %H:%M:%S UTC."
+        )
+    return (
+        f"This run's limits are {', '.join(limits[:-1])} and {limits[-1]}.{start} "
+        "Sub-agents' tool calls count toward them. Plan the work to fit, and "
+        "keep enough room to write your report."
+    )