fix: enable intra-day position tracking for sell-then-buy trades

Resolves issue where sell proceeds were not immediately available for
subsequent buy orders within the same trading session.

Problem:
- Both buy() and sell() independently queried database for starting position
- Multiple trades within same day all saw pre-trade cash balance
- Agents couldn't rebalance portfolios (sell + buy) in single session

Solution:
- ContextInjector maintains in-memory position state during trading session
- Position updates accumulate after each successful trade
- Position state injected into buy/sell via _current_position parameter
- Reset position state at start of each trading day

Changes:
- agent/context_injector.py: Add position tracking with reset_position()
- agent_tools/tool_trade.py: Accept _current_position in buy/sell functions
- agent/base_agent/base_agent.py: Reset position state daily
- tests: Add 13 comprehensive tests for position tracking

All new tests pass. Backward compatible, no schema changes required.
This commit is contained in:
2025-11-05 06:56:54 -05:00
parent 462de3adeb
commit e20dce7432
5 changed files with 449 additions and 14 deletions

View File

@@ -295,3 +295,190 @@ def test_sell_writes_to_actions_table(test_db, monkeypatch):
assert row[1] == 'AAPL'
assert row[2] == 5
assert row[3] == 160.0
def test_intraday_position_tracking_sell_then_buy(test_db, monkeypatch):
"""Test that sell proceeds are immediately available for subsequent buys."""
db, trading_day_id = test_db
# Setup: Create starting position with AAPL shares and limited cash
db.create_holding(trading_day_id, 'AAPL', 10)
db.connection.commit()
# Create a mock connection wrapper
class MockConnection:
def __init__(self, real_conn):
self.real_conn = real_conn
def cursor(self):
return self.real_conn.cursor()
def commit(self):
return self.real_conn.commit()
def rollback(self):
return self.real_conn.rollback()
def close(self):
pass
mock_conn = MockConnection(db.connection)
monkeypatch.setattr('agent_tools.tool_trade.get_db_connection',
lambda x: mock_conn)
# Mock get_current_position_from_db to return starting position
monkeypatch.setattr('agent_tools.tool_trade.get_current_position_from_db',
lambda job_id, sig, date: ({'CASH': 500.0, 'AAPL': 10}, 0))
monkeypatch.setenv('RUNTIME_ENV_PATH', '/tmp/test_runtime_intraday.json')
import json
with open('/tmp/test_runtime_intraday.json', 'w') as f:
json.dump({
'TODAY_DATE': '2025-01-15',
'SIGNATURE': 'test-model',
'JOB_ID': 'test-job-123',
'TRADING_DAY_ID': trading_day_id
}, f)
# Mock prices: AAPL sells for 200, MSFT costs 150
def mock_get_prices(date, symbols):
if 'AAPL' in symbols:
return {'AAPL_price': 200.0}
elif 'MSFT' in symbols:
return {'MSFT_price': 150.0}
return {}
monkeypatch.setattr('agent_tools.tool_trade.get_open_prices', mock_get_prices)
# Step 1: Sell 3 shares of AAPL for 600.0
# Starting cash: 500.0, proceeds: 600.0, new cash: 1100.0
result_sell = _sell_impl(
symbol='AAPL',
amount=3,
signature='test-model',
today_date='2025-01-15',
job_id='test-job-123',
trading_day_id=trading_day_id,
_current_position=None # Use database position (starting position)
)
assert 'error' not in result_sell, f"Sell should succeed: {result_sell}"
assert result_sell['CASH'] == 1100.0, "Cash should be 500 + (3 * 200) = 1100"
assert result_sell['AAPL'] == 7, "AAPL shares should be 10 - 3 = 7"
# Step 2: Buy 7 shares of MSFT for 1050.0 using the position from the sell
# This should work because we pass the updated position from step 1
result_buy = _buy_impl(
symbol='MSFT',
amount=7,
signature='test-model',
today_date='2025-01-15',
job_id='test-job-123',
trading_day_id=trading_day_id,
_current_position=result_sell # Use position from sell
)
assert 'error' not in result_buy, f"Buy should succeed with sell proceeds: {result_buy}"
assert result_buy['CASH'] == 50.0, "Cash should be 1100 - (7 * 150) = 50"
assert result_buy['MSFT'] == 7, "MSFT shares should be 7"
assert result_buy['AAPL'] == 7, "AAPL shares should still be 7"
# Verify both actions were recorded
cursor = db.connection.execute("""
SELECT action_type, symbol, quantity, price
FROM actions
WHERE trading_day_id = ?
ORDER BY created_at
""", (trading_day_id,))
actions = cursor.fetchall()
assert len(actions) == 2, "Should have 2 actions (sell + buy)"
assert actions[0][0] == 'sell' and actions[0][1] == 'AAPL'
assert actions[1][0] == 'buy' and actions[1][1] == 'MSFT'
def test_intraday_tracking_without_position_injection_fails(test_db, monkeypatch):
"""Test that without position injection, sell proceeds are NOT available for subsequent buys."""
db, trading_day_id = test_db
# Setup: Create starting position with AAPL shares and limited cash
db.create_holding(trading_day_id, 'AAPL', 10)
db.connection.commit()
# Create a mock connection wrapper
class MockConnection:
def __init__(self, real_conn):
self.real_conn = real_conn
def cursor(self):
return self.real_conn.cursor()
def commit(self):
return self.real_conn.commit()
def rollback(self):
return self.real_conn.rollback()
def close(self):
pass
mock_conn = MockConnection(db.connection)
monkeypatch.setattr('agent_tools.tool_trade.get_db_connection',
lambda x: mock_conn)
# Mock get_current_position_from_db to ALWAYS return starting position
# (simulating the old buggy behavior)
monkeypatch.setattr('agent_tools.tool_trade.get_current_position_from_db',
lambda job_id, sig, date: ({'CASH': 500.0, 'AAPL': 10}, 0))
monkeypatch.setenv('RUNTIME_ENV_PATH', '/tmp/test_runtime_no_injection.json')
import json
with open('/tmp/test_runtime_no_injection.json', 'w') as f:
json.dump({
'TODAY_DATE': '2025-01-15',
'SIGNATURE': 'test-model',
'JOB_ID': 'test-job-123',
'TRADING_DAY_ID': trading_day_id
}, f)
# Mock prices
def mock_get_prices(date, symbols):
if 'AAPL' in symbols:
return {'AAPL_price': 200.0}
elif 'MSFT' in symbols:
return {'MSFT_price': 150.0}
return {}
monkeypatch.setattr('agent_tools.tool_trade.get_open_prices', mock_get_prices)
# Step 1: Sell 3 shares of AAPL
result_sell = _sell_impl(
symbol='AAPL',
amount=3,
signature='test-model',
today_date='2025-01-15',
job_id='test-job-123',
trading_day_id=trading_day_id,
_current_position=None # Don't inject position (old behavior)
)
assert 'error' not in result_sell, "Sell should succeed"
# Step 2: Try to buy 7 shares of MSFT WITHOUT passing updated position
# This should FAIL because it will query the database and get the original 500.0 cash
result_buy = _buy_impl(
symbol='MSFT',
amount=7,
signature='test-model',
today_date='2025-01-15',
job_id='test-job-123',
trading_day_id=trading_day_id,
_current_position=None # Don't inject position (old behavior)
)
# This should fail with insufficient cash
assert 'error' in result_buy, "Buy should fail without position injection"
assert result_buy['error'] == 'Insufficient cash', f"Expected insufficient cash error, got: {result_buy}"
assert result_buy['cash_available'] == 500.0, "Should see original cash, not updated cash"