API endpoint for task cost (#5377)

This adds a new API endpoint for fetching the cost of a particular task:
`tasks/<task-id>/cost`.

The feature is wired though the SDK to the MCP server. The CC system
prompt encourages the agent to use this tool instead of estimating
costs. I made some other adjustments to the system prompt to
disambiguate the user's account balance from the conversation budget.

Sourced from commit 58494fdf34b17e2e639777ab6c34d701c4f69040

Author: rgambee

futuresearch-mcp/manifest.json

@@ -92,6 +92,10 @@
       "name": "futuresearch_balance",
       "description": "Check the current billing balance for the authenticated user."
     },
+    {
+      "name": "futuresearch_task_cost",
+      "description": "Get the billed cost of a completed task."
+    },
     {
       "name": "futuresearch_browse_lists",
       "description": "Browse available reference lists of well-known entities."

futuresearch-mcp/src/futuresearch_mcp/models.py

@@ -637,6 +637,19 @@ def validate_task_id(cls, v: str) -> str:
         return _validate_task_id(v)
 
 
+class TaskCostInput(BaseModel):
+    """Input for getting the billed cost of a task."""
+
+    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
+
+    task_id: str = Field(..., description="The task ID to check the cost for.")
+
+    @field_validator("task_id")
+    @classmethod
+    def validate_task_id(cls, v: str) -> str:
+        return _validate_task_id(v)
+
+
 def _validate_output_path(v: str | None) -> str | None:
     """Validate output_path ends in .csv and parent directory exists."""
     if v is not None:

futuresearch-mcp/src/futuresearch_mcp/tools.py

@@ -15,6 +15,7 @@
 from futuresearch.constants import FuturesearchError as EveryrowError
 from futuresearch.generated.api.billing import get_billing_balance_billing_get
 from futuresearch.generated.api.tasks import get_task_status_tasks_task_id_status_get
+from futuresearch.generated.models.task_cost_status import TaskCostStatus
 from futuresearch.generated.models.task_status import TaskStatus
 from futuresearch.ops import (
     _submit_agent_map,
@@ -27,7 +28,7 @@
     merge_async,
 )
 from futuresearch.session import list_sessions
-from futuresearch.task import cancel_task
+from futuresearch.task import cancel_task, get_task_cost
 from mcp.types import CallToolResult, TextContent, ToolAnnotations
 from pydantic import BaseModel, create_model
 
@@ -49,6 +50,7 @@
     RankInput,
     SingleAgentInput,
     StdioResultsInput,
+    TaskCostInput,
     UploadDataInput,
     UseListInput,
 )
@@ -1337,6 +1339,58 @@ async def futuresearch_balance(ctx: FuturesearchContext) -> list[TextContent]:
     ]
 
 
+@mcp.tool(
+    name="futuresearch_task_cost",
+    structured_output=False,
+    annotations=ToolAnnotations(
+        title="Get Task Cost",
+        readOnlyHint=True,
+        destructiveHint=False,
+        idempotentHint=True,
+        openWorldHint=False,
+    ),
+)
+async def futuresearch_task_cost(
+    params: TaskCostInput, ctx: FuturesearchContext
+) -> list[TextContent]:
+    """Get the billed cost of a completed task.
+
+    Returns the amount charged in dollars. There is a delay between task completion and
+    cost calculation. Returns 'pending' if the cost hasn't settled yet.
+    """
+    logger.info("futuresearch_task_cost: task_id=%s", params.task_id)
+    client = _get_client(ctx)
+
+    try:
+        response = await get_task_cost(
+            task_id=UUID(params.task_id),
+            client=client,
+        )
+    except Exception:
+        logger.exception("Failed to get task cost for task %s", params.task_id)
+        return [
+            TextContent(
+                type="text",
+                text=f"Error retrieving cost for task {params.task_id}. Please try again.",
+            )
+        ]
+
+    if response.status == TaskCostStatus.PENDING:
+        return [
+            TextContent(
+                type="text",
+                text=f"Cost for task {params.task_id} is still being calculated. Try again in ~30 seconds.",
+            )
+        ]
+
+    return [
+        TextContent(
+            type="text",
+            text=f"Task {params.task_id} cost: ${response.cost_dollars:.2f}",
+        )
+    ]
+
+
 @mcp.tool(
     name="futuresearch_list_session_tasks",
     structured_output=False,

futuresearch-mcp/tests/test_mcp_e2e.py

@@ -168,9 +168,10 @@ async def test_list_tools(self, _http_state):
                     "futuresearch_merge",
                     "futuresearch_progress",
                     "futuresearch_rank",
-                    "futuresearch_status",
                     "futuresearch_results",
                     "futuresearch_single_agent",
+                    "futuresearch_status",
+                    "futuresearch_task_cost",
                     "futuresearch_upload_data",
                     "futuresearch_use_list",
                 ]

src/futuresearch/generated/api/tasks/get_task_cost_tasks_task_id_cost_get.py

@@ -0,0 +1,179 @@
+from http import HTTPStatus
+from typing import Any
+from urllib.parse import quote
+from uuid import UUID
+
+import httpx
+
+from ... import errors
+from ...client import AuthenticatedClient, Client
+from ...models.error_response import ErrorResponse
+from ...models.http_validation_error import HTTPValidationError
+from ...models.task_cost_response import TaskCostResponse
+from ...types import Response
+
+
+def _get_kwargs(
+    task_id: UUID,
+) -> dict[str, Any]:
+    _kwargs: dict[str, Any] = {
+        "method": "get",
+        "url": "/tasks/{task_id}/cost".format(
+            task_id=quote(str(task_id), safe=""),
+        ),
+    }
+
+    return _kwargs
+
+
+def _parse_response(
+    *, client: AuthenticatedClient | Client, response: httpx.Response
+) -> ErrorResponse | HTTPValidationError | TaskCostResponse | None:
+    if response.status_code == 200:
+        response_200 = TaskCostResponse.from_dict(response.json())
+
+        return response_200
+
+    if response.status_code == 404:
+        response_404 = ErrorResponse.from_dict(response.json())
+
+        return response_404
+
+    if response.status_code == 422:
+        response_422 = HTTPValidationError.from_dict(response.json())
+
+        return response_422
+
+    if client.raise_on_unexpected_status:
+        raise errors.UnexpectedStatus(response.status_code, response.content)
+    else:
+        return None
+
+
+def _build_response(
+    *, client: AuthenticatedClient | Client, response: httpx.Response
+) -> Response[ErrorResponse | HTTPValidationError | TaskCostResponse]:
+    return Response(
+        status_code=HTTPStatus(response.status_code),
+        content=response.content,
+        headers=response.headers,
+        parsed=_parse_response(client=client, response=response),
+    )
+
+
+def sync_detailed(
+    task_id: UUID,
+    *,
+    client: AuthenticatedClient,
+) -> Response[ErrorResponse | HTTPValidationError | TaskCostResponse]:
+    """Get task cost
+
+     Get the billed cost of a task. Returns status 'pending' until the cost has been calculated, then
+    'settled' with the final cost in dollars.
+
+    Args:
+        task_id (UUID):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[ErrorResponse | HTTPValidationError | TaskCostResponse]
+    """
+
+    kwargs = _get_kwargs(
+        task_id=task_id,
+    )
+
+    response = client.get_httpx_client().request(
+        **kwargs,
+    )
+
+    return _build_response(client=client, response=response)
+
+
+def sync(
+    task_id: UUID,
+    *,
+    client: AuthenticatedClient,
+) -> ErrorResponse | HTTPValidationError | TaskCostResponse | None:
+    """Get task cost
+
+     Get the billed cost of a task. Returns status 'pending' until the cost has been calculated, then
+    'settled' with the final cost in dollars.
+
+    Args:
+        task_id (UUID):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        ErrorResponse | HTTPValidationError | TaskCostResponse
+    """
+
+    return sync_detailed(
+        task_id=task_id,
+        client=client,
+    ).parsed
+
+
+async def asyncio_detailed(
+    task_id: UUID,
+    *,
+    client: AuthenticatedClient,
+) -> Response[ErrorResponse | HTTPValidationError | TaskCostResponse]:
+    """Get task cost
+
+     Get the billed cost of a task. Returns status 'pending' until the cost has been calculated, then
+    'settled' with the final cost in dollars.
+
+    Args:
+        task_id (UUID):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        Response[ErrorResponse | HTTPValidationError | TaskCostResponse]
+    """
+
+    kwargs = _get_kwargs(
+        task_id=task_id,
+    )
+
+    response = await client.get_async_httpx_client().request(**kwargs)
+
+    return _build_response(client=client, response=response)
+
+
+async def asyncio(
+    task_id: UUID,
+    *,
+    client: AuthenticatedClient,
+) -> ErrorResponse | HTTPValidationError | TaskCostResponse | None:
+    """Get task cost
+
+     Get the billed cost of a task. Returns status 'pending' until the cost has been calculated, then
+    'settled' with the final cost in dollars.
+
+    Args:
+        task_id (UUID):
+
+    Raises:
+        errors.UnexpectedStatus: If the server returns an undocumented status code and Client.raise_on_unexpected_status is True.
+        httpx.TimeoutException: If the request takes longer than Client.timeout.
+
+    Returns:
+        ErrorResponse | HTTPValidationError | TaskCostResponse
+    """
+
+    return (
+        await asyncio_detailed(
+            task_id=task_id,
+            client=client,
+        )
+    ).parsed

src/futuresearch/generated/models/__init__.py

@@ -63,6 +63,8 @@
 from .subscription_info import SubscriptionInfo
 from .subscription_status_response import SubscriptionStatusResponse
 from .subscription_tier import SubscriptionTier
+from .task_cost_response import TaskCostResponse
+from .task_cost_status import TaskCostStatus
 from .task_progress_info import TaskProgressInfo
 from .task_result_response import TaskResultResponse
 from .task_result_response_data_type_0_item import TaskResultResponseDataType0Item
@@ -146,6 +148,8 @@
     "SubscriptionInfo",
     "SubscriptionStatusResponse",
     "SubscriptionTier",
+    "TaskCostResponse",
+    "TaskCostStatus",
     "TaskProgressInfo",
     "TaskResultResponse",
     "TaskResultResponseDataType0Item",

src/futuresearch/generated/models/forecast_operation.py

@@ -29,7 +29,8 @@ class ForecastOperation:
             question/scenario to forecast.
         forecast_type (ForecastOperationForecastType): Type of forecast. 'binary': yes/no probability (0-100) for
             questions like 'Will X happen?'. 'numeric': percentile estimates (p10-p90) for questions like 'What will the
-            price/value/count be?'. Requires output_field when 'numeric'.
+            price/value/count be?'. 'date': date percentile estimates (p10-p90) as YYYY-MM-DD strings for timing questions
+            like 'When will X happen?'. Requires output_field when 'numeric' or 'date'.
         session_id (None | Unset | UUID): Session ID. If not provided, a new session is auto-created for this task.
         webhook_url (None | str | Unset): Optional URL to receive a POST callback when the task completes or fails.
         output_field (None | str | Unset): Name of the numeric quantity being forecast (e.g. 'price', 'count'). Required

src/futuresearch/generated/models/forecast_operation_forecast_type.py

@@ -3,8 +3,8 @@
 
 class ForecastOperationForecastType(str, Enum):
     BINARY = "binary"
-    NUMERIC = "numeric"
     DATE = "date"
+    NUMERIC = "numeric"
 
     def __str__(self) -> str:
         return str(self.value)

src/futuresearch/generated/models/progress_summary_entry.py

@@ -21,14 +21,16 @@ class ProgressSummaryEntry:
         iteration_number (int): Iteration within the trace
         summary (str): LLM-generated progress summary text
         updated_at (str): When this summary was created
-        row_index (int | None | Unset): Input row index this trace covers (if resolved)
+        row_indices (list[int] | Unset): Input row indices this trace covers
+        row_index (int | None | Unset): Deprecated: use row_indices instead
     """
 
     trace_id: UUID
     task_id: UUID
     iteration_number: int
     summary: str
     updated_at: str
+    row_indices: list[int] | Unset = UNSET
     row_index: int | None | Unset = UNSET
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
@@ -43,6 +45,10 @@ def to_dict(self) -> dict[str, Any]:
 
         updated_at = self.updated_at
 
+        row_indices: list[int] | Unset = UNSET
+        if not isinstance(self.row_indices, Unset):
+            row_indices = self.row_indices
+
         row_index: int | None | Unset
         if isinstance(self.row_index, Unset):
             row_index = UNSET
@@ -60,6 +66,8 @@ def to_dict(self) -> dict[str, Any]:
                 "updated_at": updated_at,
             }
         )
+        if row_indices is not UNSET:
+            field_dict["row_indices"] = row_indices
         if row_index is not UNSET:
             field_dict["row_index"] = row_index
 
@@ -78,6 +86,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
 
         updated_at = d.pop("updated_at")
 
+        row_indices = cast(list[int], d.pop("row_indices", UNSET))
+
         def _parse_row_index(data: object) -> int | None | Unset:
             if data is None:
                 return data
@@ -93,6 +103,7 @@ def _parse_row_index(data: object) -> int | None | Unset:
             iteration_number=iteration_number,
             summary=summary,
             updated_at=updated_at,
+            row_indices=row_indices,
             row_index=row_index,
         )