Merge pull request modelcontextprotocol#21 from jkoelker/jk/discord

feat(tools): surface order metadata from placements

Author: jkoelker

src/schwab_mcp/tools/orders.py

@@ -5,6 +5,11 @@
 import copy
 import datetime
 from mcp.server.fastmcp import FastMCP
+from schwab.utils import (
+    AccountHashMismatchException,
+    UnsuccessfulOrderException,
+    Utils as SchwabUtils,
+)
 from schwab.orders.common import first_triggers_second as trigger_builder
 from schwab.orders.common import one_cancels_other as oco_builder
 from schwab.orders.options import OptionSymbol
@@ -30,7 +35,7 @@
     option_sell_to_open_limit,
     option_sell_to_open_market,
 )
-from schwab_mcp.tools.utils import JSONType, call
+from schwab_mcp.tools.utils import JSONType, ResponseHandler, call
 
 
 # Internal helper function to apply session and duration settings
@@ -168,6 +173,33 @@ def _build_option_order_spec(
         )
 
 
+def _order_response_handler(ctx: SchwabContext, account_hash: str) -> ResponseHandler:
+    utils = SchwabUtils(ctx.client, account_hash)
+
+    def handler(response: Any) -> tuple[bool, JSONType]:
+        headers = getattr(response, "headers", {})
+        location = headers.get("Location") if headers else None
+
+        try:
+            order_id = utils.extract_order_id(response)
+        except (AccountHashMismatchException, UnsuccessfulOrderException):
+            order_id = None
+
+        if order_id is None and location is None:
+            return False, None
+
+        payload: dict[str, Any] = {}
+        if order_id is not None:
+            payload["orderId"] = order_id
+            payload["accountHash"] = account_hash
+        if location is not None:
+            payload["location"] = location
+
+        return True, payload
+
+    return handler
+
+
 async def get_order(
     ctx: SchwabContext,
     account_hash: Annotated[str, "Account hash for the Schwab account"],
@@ -286,7 +318,10 @@ async def place_equity_order(
 
     # Place the order
     return await call(
-        client.place_order, account_hash=account_hash, order_spec=order_spec_dict
+        client.place_order,
+        account_hash=account_hash,
+        order_spec=order_spec_dict,
+        response_handler=_order_response_handler(ctx, account_hash),
     )
 
 
@@ -332,7 +367,10 @@ async def place_option_order(
 
     # Place the order
     return await call(
-        client.place_order, account_hash=account_hash, order_spec=order_spec_dict
+        client.place_order,
+        account_hash=account_hash,
+        order_spec=order_spec_dict,
+        response_handler=_order_response_handler(ctx, account_hash),
     )
 
 
@@ -433,7 +471,10 @@ async def place_one_cancels_other_order(
     client = ctx.orders
 
     return await call(
-        client.place_order, account_hash=account_hash, order_spec=oco_order_spec
+        client.place_order,
+        account_hash=account_hash,
+        order_spec=oco_order_spec,
+        response_handler=_order_response_handler(ctx, account_hash),
     )
 
 
@@ -491,6 +532,7 @@ async def place_first_triggers_second_order(
         client.place_order,
         account_hash=account_hash,
         order_spec=trigger_order_dict,
+        response_handler=_order_response_handler(ctx, account_hash),
     )
 
 
@@ -604,6 +646,7 @@ async def place_bracket_order(
         client.place_order,
         account_hash=account_hash,
         order_spec=bracket_order_dict,
+        response_handler=_order_response_handler(ctx, account_hash),
     )
 
 
@@ -669,7 +712,10 @@ async def place_option_combo_order(
         )
 
     return await call(
-        ctx.orders.place_order, account_hash=account_hash, order_spec=builder.build()
+        ctx.orders.place_order,
+        account_hash=account_hash,
+        order_spec=builder.build(),
+        response_handler=_order_response_handler(ctx, account_hash),
     )
 
 

src/schwab_mcp/tools/utils.py

@@ -8,6 +8,9 @@
 JSONType: TypeAlias = JSONPrimitive | dict[str, Any] | list[Any]
 
 
+ResponseHandler: TypeAlias = Callable[[Any], tuple[bool, JSONType]]
+
+
 class SchwabAPIError(Exception):
     """Represents an error response returned from the Schwab API."""
 
@@ -23,8 +26,18 @@ def __init__(
         )
 
 
-async def call(func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> JSONType:
-    """Call a Schwab client endpoint and return its JSON payload."""
+async def call(
+    func: Callable[..., Awaitable[Any]],
+    *args: Any,
+    response_handler: ResponseHandler | None = None,
+    **kwargs: Any,
+) -> JSONType:
+    """Call a Schwab client endpoint and return its JSON payload.
+
+    When ``response_handler`` is provided, it can opt to handle the response
+    by returning ``(True, payload)``. Returning ``(False, _)`` delegates back to
+    the default JSON parsing behavior.
+    """
 
     response = await func(*args, **kwargs)
     try:
@@ -36,6 +49,11 @@ async def call(func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -
             body=response.text,
         ) from exc
 
+    if response_handler is not None:
+        handled, payload = response_handler(response)
+        if handled:
+            return payload
+
     # Handle responses with no content
     # 204 No Content: explicit no-content response
     # 201 Created: order placement endpoints return empty body with Location header
@@ -55,4 +73,4 @@ async def call(func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -
         raise ValueError("Expected JSON response from Schwab endpoint") from exc
 
 
-__all__ = ["call", "JSONType", "SchwabAPIError"]
+__all__ = ["call", "JSONType", "SchwabAPIError", "ResponseHandler"]

tests/test_orders.py

@@ -116,3 +116,55 @@ async def fake_call(func, *args, **kwargs):
         client.Order.Status.FILLED,
         client.Order.Status.CANCELED,
     ]
+
+
+def test_place_equity_order_returns_order_metadata():
+    account_hash = "abc123"
+    order_id = 987654321
+    location = (
+        f"https://api.schwabapi.com/trader/v1/accounts/{account_hash}/orders/{order_id}"
+    )
+
+    class DummyResponse:
+        status_code = 201
+        url = f"https://api.schwabapi.com/trader/v1/accounts/{account_hash}/orders"
+        text = ""
+        content = b""
+        headers = {"Location": location}
+        is_error = False
+
+        def raise_for_status(self) -> None:
+            return None
+
+    class DummyPlaceOrderClient(DummyOrdersClient):
+        def __init__(self) -> None:
+            self.captured: dict[str, Any] | None = None
+
+        async def place_order(self, *args, **kwargs):
+            self.captured = {"args": args, "kwargs": kwargs}
+            return DummyResponse()
+
+    client = DummyPlaceOrderClient()
+    ctx = make_ctx(client)
+
+    result = run(
+        orders.place_equity_order(
+            ctx,
+            account_hash,
+            "SPY",
+            1,
+            "buy",
+            "market",
+        )
+    )
+
+    assert result == {
+        "orderId": order_id,
+        "accountHash": account_hash,
+        "location": location,
+    }
+
+    captured = client.captured
+    assert captured is not None
+    assert captured["kwargs"]["account_hash"] == account_hash
+    assert captured["kwargs"]["order_spec"]["orderStrategyType"] == "SINGLE"