"""Review fixes: partial-sell realized P&L, training reset leaves the circuit armed, close guard + runner lock, confirmation polling past lastValidBlockHeight.""" import threading import uuid import pytest from fly_trader import config from fly_trader.agent import runner from fly_trader.db.connection import transaction from fly_trader.execution import broker_live, ledger from fly_trader.execution.broker_live import FillResult from fly_trader.execution.broker_paper import PaperBroker from fly_trader.execution.worker import ExecRequest, ExecutionWorker def _pos(conn, pid): return conn.execute("SELECT * FROM positions WHERE id=%s", (pid,)).fetchone() # ---- 1. partial sells accumulate realized P&L; paper cash stays consistent ---- def test_paper_partial_then_close_accumulates_realized(db_conn): book, mint = f"replay_t{uuid.uuid4().hex[:8]}", f"M{uuid.uuid4().hex[:10]}" pid = ledger.open_position(db_conn, book=book, mint=mint, pool=None, qty_raw=1_000_000, cost_sol=1.0, entry_price=1.0, decision_id=None, fees_sol=0.0) assert ledger.paper_cash(db_conn, book) == pytest.approx(config.CAPITAL_SOL - 1.0) pb = PaperBroker(book) kw = dict(decision_id=None, res_quote_sol=100.0, mcap_sol=100.0, program_label=None, forced_kind=None) f1 = pb.sell(db_conn, position=ledger.open_positions(db_conn, book)[0], price=1.2, fraction=0.5, **kw) assert f1.ok and 0 < f1.sol_delta < 0.6 p = _pos(db_conn, pid) assert p["status"] == "open" and int(p["qty"]) == 500_000 and p["cost_sol"] == pytest.approx(0.5) assert p["realized_sol"] == pytest.approx(f1.sol_delta - 0.5) assert ledger.paper_cash(db_conn, book) == pytest.approx(config.CAPITAL_SOL - 1.0 + f1.sol_delta) stale = ledger.open_positions(db_conn, book)[0] f2 = pb.sell(db_conn, position=stale, price=0.8, **kw) assert f2.ok and "realized" in f2.reason p = _pos(db_conn, pid) assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(f1.sol_delta + f2.sol_delta - 1.0) assert ledger.paper_cash(db_conn, book) == pytest.approx(config.CAPITAL_SOL - 1.0 + f1.sol_delta + f2.sol_delta) # 4. the close guard: a second close of the same position books nothing and records no fill f3 = pb.sell(db_conn, position=stale, price=0.8, **kw) assert not f3.ok and f3.fill_id is None assert ledger.close_position(db_conn, position_id=pid, exit_price=1.0, proceeds_sol=5.0, fees_sol=0.0, decision_id=None, forced_kind=None) is None assert _pos(db_conn, pid)["realized_sol"] == pytest.approx(f1.sol_delta + f2.sol_delta - 1.0) n = db_conn.execute("SELECT count(*) AS n FROM fills WHERE book=%s AND mint=%s AND side='sell'", (book, mint)).fetchone()["n"] assert n == 2 class _FakeBroker: def __init__(self): self.next: FillResult | None = None self.kw: dict = {} def swap(self, conn, **kw): self.kw = kw return self.next def test_worker_partial_then_close_books_each_sale_once(): mint = f"M{uuid.uuid4().hex[:10]}" with transaction() as conn: circ = dict(conn.execute("SELECT fail_count, tripped FROM circuit_state WHERE id=1").fetchone()) pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=1.0, entry_price=1.0, decision_id=None, fees_sol=0.0) fb = _FakeBroker() w = ExecutionWorker(fb) try: req = ExecRequest(decision_id=None, mint=mint, pool=None, side="sell", amount_in=400_000, slippage_bps=100, max_slippage_bps=100, decimals=6, position_id=pid) fb.next = FillResult(True, None, None, "s1", "confirmed", 0, None, -400_000, 500_000_000, 1.25, 1) r1 = w._execute(req) assert r1.ok and r1.realized_sol == pytest.approx(0.1) fb.next = FillResult(True, None, None, "s2", "confirmed", 0, None, -600_000, 900_000_000, 1.5, 1) r2 = w._execute(req) assert r2.ok and r2.realized_sol == pytest.approx(0.3) fb.next = FillResult(True, None, None, "s3", "confirmed", 0, None, -1, 1, 1.0, 1) r3 = w._execute(req) assert r3.ok and r3.realized_sol is None with transaction() as conn: p = _pos(conn, pid) assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(0.4) and p["cost_sol"] == pytest.approx(0.6) finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE id=%s", (pid,)) conn.execute("UPDATE circuit_state SET fail_count=%s, tripped=%s WHERE id=1", (circ["fail_count"], circ["tripped"])) # ---- 2. a training reset never re-arms the circuit ---- def test_training_reset_leaves_circuit_alone(tmp_path, monkeypatch): from fly_trader.ops.reset import reset_training_state monkeypatch.setattr(config, "PG_ARCHIVE_DIR", tmp_path / "arch") monkeypatch.setattr(config, "BRAIN_DIR", tmp_path / "brain") cols = "fail_count, tripped, kill_switch, kill_reason, peak_wealth" with transaction() as conn: before = dict(conn.execute(f"SELECT {cols} FROM circuit_state WHERE id=1").fetchone()) conn.execute("UPDATE circuit_state SET fail_count=3, tripped=true, kill_switch=true, kill_reason='dd', peak_wealth=5.0 WHERE id=1") try: reset_training_state(reason="test", archive=False) with transaction() as conn: after = dict(conn.execute(f"SELECT {cols} FROM circuit_state WHERE id=1").fetchone()) assert after == {"fail_count": 3, "tripped": True, "kill_switch": True, "kill_reason": "dd", "peak_wealth": 5.0} finally: with transaction() as conn: conn.execute("UPDATE circuit_state SET fail_count=%(fail_count)s, tripped=%(tripped)s, kill_switch=%(kill_switch)s, " "kill_reason=%(kill_reason)s, peak_wealth=%(peak_wealth)s WHERE id=1", before) # ---- 4. one runner per database ---- def test_runner_lock_refuses_a_second_runner(monkeypatch): monkeypatch.setattr(runner, "setup", lambda name: None) ran = [] monkeypatch.setattr(runner, "_main", lambda stop_event=None: ran.append(1)) held = runner.acquire_runner_lock() try: with pytest.raises(RuntimeError, match="runner lock"): runner.acquire_runner_lock() with pytest.raises(RuntimeError, match="runner lock"): runner.main() assert ran == [] finally: held.close() runner.main() # lock free again; main releases it on return assert ran == [1] runner.acquire_runner_lock().close() # ---- 5. a processed signature is polled past lastValidBlockHeight ---- class _StatusRpc: def __init__(self, statuses): self.statuses = list(statuses) def get_signature_statuses(self, sigs): return [self.statuses.pop(0)] def get_block_height(self): return 10_000 def test_await_confirmation_keeps_polling_a_processed_signature(monkeypatch): monkeypatch.setattr(broker_live.time, "sleep", lambda s: None) proc = {"err": None, "confirmationStatus": "processed"} rpc = _StatusRpc([proc, proc, {"err": None, "confirmationStatus": "confirmed"}]) assert broker_live.await_confirmation(rpc, "sig", 100)[0] == "confirmed" rpc = _StatusRpc([proc, None]) assert broker_live.await_confirmation(rpc, "sig", 100)[0] == "expired" # ---- 6. live keeps up with paper: retryable failures, sweeps, a buy behind a sale, rent back, a clean stop ---- class _ScriptBroker: """Fills from a script, in order; ``gate`` holds each swap until set (an order in flight); every call recorded.""" def __init__(self, fills, gate=None): self.fills, self.gate = list(fills), gate self.calls, self.started = [], threading.Event() def swap(self, conn, **kw): self.started.set(); self.calls.append(kw) if self.gate is not None: self.gate.wait(10) return self.fills.pop(0) @pytest.fixture def circuit1(): with transaction() as conn: before = dict(conn.execute("SELECT fail_count, tripped, last_failure_ts FROM circuit_state WHERE id=1").fetchone()) conn.execute("UPDATE circuit_state SET fail_count = 0, tripped = false WHERE id = 1") yield with transaction() as conn: conn.execute("UPDATE circuit_state SET fail_count=%(fail_count)s, tripped=%(tripped)s, last_failure_ts=%(last_failure_ts)s WHERE id=1", before) def _req(mint, side, amount_in=1_000_000, **kw): return ExecRequest(decision_id=None, mint=mint, pool=None, side=side, amount_in=amount_in, slippage_bps=300, max_slippage_bps=500, decimals=6, **kw) def _fails(): with transaction() as conn: return conn.execute("SELECT fail_count FROM circuit_state WHERE id=1").fetchone()["fail_count"] def _fill(token_delta, lamports_delta, price=1e-7): return FillResult(True, None, None, "sig", "confirmed", 0, None, token_delta, lamports_delta, price, 1) def test_failed_buys_that_moved_nothing_may_go_again_and_a_price_that_ran_is_no_execution_failure(circuit1): mint = f"M{uuid.uuid4().hex[:10]}" fb = _FakeBroker(); w = ExecutionWorker(fb) try: fb.next = FillResult(False, 1, None, None, "price_ran", None, "price ran: quoted 1.1e-07 SOL per token, +10.0% over the decision's 1e-07", 0, 0, None, 1) r = w._execute(_req(mint, "buy", 10 ** 8, ref_price=1e-7)) assert not r.ok and r.retryable and r.error.startswith("price ran") and (fb.kw["ref_price"], fb.kw["decimals"]) == (1e-7, 6) and _fails() == 0 fb.next = FillResult(False, 2, None, None, "order_error", 1, "Insufficient funds", 0, 0, None, 1) r = w._execute(_req(mint, "buy", 10 ** 8)) assert not r.ok and r.retryable and _fails() == 1 # a failure, but nothing landed fb.next = FillResult(False, 3, 99, "sig", "confirmed", 0, None, 0, 0, None, 1) assert not w._execute(_req(mint, "buy", 10 ** 8)).retryable and _fails() == 2 # an unverified fill row: never sent again with transaction() as conn: ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.1, entry_price=1e-7, decision_id=None, fees_sol=0.0) fb.kw = {} r = w._execute(_req(mint, "buy", 10 ** 8)) # the sale queued before it left the position open assert not r.ok and r.retryable and r.error.startswith("held") and fb.kw == {} and _fails() == 2 finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) def test_a_dead_bag_sweep_books_its_proceeds_onto_the_written_off_position(circuit1): mint = f"M{uuid.uuid4().hex[:10]}" with transaction() as conn: pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.2, entry_price=2e-7, decision_id=None, fees_sol=0.0) ledger.close_position(conn, position_id=pid, exit_price=0.0, proceeds_sol=0.0, fees_sol=0.0, decision_id=None, forced_kind="dead_bag") fb = _FakeBroker(); w = ExecutionWorker(fb) try: fb.next = _fill(-1_000_000, 30_000_000, 3e-8) r = w._execute(_req(mint, "sell", forced_kind="dead_bag_sweep")) assert r.ok and r.position_id == pid and r.closed and r.realized_sol == pytest.approx(0.03) with transaction() as conn: p = _pos(conn, pid) assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(-0.2 + 0.03) finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) def test_a_sale_that_empties_its_mint_leaves_the_token_accounts_close_to_the_mirror(circuit1): """Live money first: the worker goes on to the next order; the mirror closes the emptied account mid-minute.""" mint = f"M{uuid.uuid4().hex[:10]}" with transaction() as conn: pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.2, entry_price=2e-7, decision_id=None, fees_sol=0.0) fb = _ScriptBroker([_fill(-1_000_000, 250_000_000, 2.5e-7)]); w = ExecutionWorker(fb) try: assert w.submit(_req(mint, "sell", position_id=pid)) r = w.results.get(timeout=10) assert r.ok and r.closed and w.emptied == {mint} and r.realized_sol == pytest.approx(0.05) with transaction() as conn: p = _pos(conn, pid) assert p["status"] == "closed" and p["realized_sol"] == pytest.approx(0.25 - 0.2) and w.pending() == set() finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) def test_a_first_buys_token_account_rent_is_a_deposit_its_close_returns(circuit1): """The rent a first buy pays for the mint's token account is kept apart from the position's cost (a live buy costs what its paper twin does) and counted in live wealth while held; the sale's close returns it, and only the close's fee is the position's.""" mint = f"M{uuid.uuid4().hex[:10]}"; rent, fee = 1_513_840, 5_000 buy = FillResult(True, None, None, "sig", "confirmed", 0, None, 1_000_000, -200_000_000, 2e-7, 1, rent) fb = _ScriptBroker([buy, _fill(-1_000_000, 250_000_000, 2.5e-7)]); w = ExecutionWorker(fb) try: assert w.submit(_req(mint, "buy", 200_000_000)) r1 = w.results.get(timeout=10) with transaction() as conn: p = _pos(conn, r1.position_id); held = ledger.deposits_held(conn, "live") assert r1.ok and p["cost_sol"] == pytest.approx(0.2) and p["deposit_sol"] == pytest.approx(rent / 1e9) and held >= rent / 1e9 assert w.submit(_req(mint, "sell", position_id=r1.position_id)) r2 = w.results.get(timeout=10) with transaction() as conn: before = ledger.deposits_held(conn, "live") booked = ledger.book_account_close(conn, "live", mint, rent - fee) # the mirror's close (agent/fly_live.close_accounts) p = _pos(conn, r1.position_id); after = ledger.deposits_held(conn, "live") assert r2.closed and w.emptied == {mint} and r2.realized_sol == pytest.approx(0.05) assert booked == (pytest.approx(-fee / 1e9), r1.position_id) and before - after == pytest.approx(rent / 1e9) assert p["deposit_refunded"] and p["realized_sol"] == pytest.approx(0.25 - 0.2 - fee / 1e9) finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) def test_a_multi_hop_buy_has_the_mirror_close_its_intermediate_token_accounts(circuit1): """The accounts a multi-hop route opened for its intermediate tokens and left empty are the mirror's to close mid-minute, as a sale's are: their rent comes back. It is neither the position's cost nor its deposit.""" mint, hop = f"M{uuid.uuid4().hex[:10]}", f"H{uuid.uuid4().hex[:10]}" buy = FillResult(True, None, None, "sig", "confirmed", 0, None, 1_000_000, -200_000_000, 2e-7, 1, 0, {hop: 1_488_440}) fb = _ScriptBroker([buy]); w = ExecutionWorker(fb) try: assert w.submit(_req(mint, "buy", 200_000_000)) r = w.results.get(timeout=10) with transaction() as conn: p = _pos(conn, r.position_id) assert r.ok and w.emptied == {hop} and p["cost_sol"] == pytest.approx(0.2) and p["deposit_sol"] is None finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) def test_a_buy_queues_behind_the_sale_of_its_mint_which_keeps_its_token_account(circuit1): """Paper sold the token and bought it back in one minute: the live buy waits behind the sale, runs once the sale has closed the position, and the token account stays for it.""" mint = f"M{uuid.uuid4().hex[:10]}" with transaction() as conn: pid = ledger.open_position(conn, book="live", mint=mint, pool=None, qty_raw=1_000_000, cost_sol=0.2, entry_price=2e-7, decision_id=None, fees_sol=0.0) gate = threading.Event() fb = _ScriptBroker([_fill(-1_000_000, 250_000_000, 2.5e-7), _fill(900_000, -200_000_000, 2.2e-7)], gate=gate); w = ExecutionWorker(fb) try: assert w.submit(_req(mint, "sell", position_id=pid)) and fb.started.wait(10) and w.selling() == {mint} assert w.submit(_req(mint, "buy", 200_000_000)) assert not w.submit(_req(mint, "buy", 1)) and not w.submit(_req(mint, "sell")) # one buy behind one sale, nothing more assert w.selling() == set() and w.pending() == {mint} gate.set() r1, r2 = w.results.get(timeout=10), w.results.get(timeout=10) assert (r1.request.side, r1.ok, r1.closed, r2.request.side, r2.ok) == ("sell", True, True, "buy", True) and w.emptied == set() with transaction() as conn: rows = conn.execute("SELECT status, qty FROM positions WHERE book = 'live' AND mint = %s ORDER BY id", (mint,)).fetchall() assert [(r["status"], int(r["qty"])) for r in rows] == [("closed", 1_000_000), ("open", 900_000)] finally: gate.set(); w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,)) def test_stop_waits_for_the_order_in_flight_and_drops_the_queued_ones(circuit1): a, b = f"M{uuid.uuid4().hex[:10]}", f"M{uuid.uuid4().hex[:10]}" gate = threading.Event() fb = _ScriptBroker([_fill(-5, 5), _fill(-5, 5)], gate=gate); w = ExecutionWorker(fb, threads=1) # one thread: b waits in the queue assert w.submit(_req(a, "sell")) and fb.started.wait(10) and w.submit(_req(b, "sell")) threading.Timer(0.3, gate.set).start() w.stop(timeout=10) assert not w.thread.is_alive() and [c["mint"] for c in fb.calls] == [a] and w.pending() == set() assert w.results.get_nowait().request.mint == a and not w.submit(_req(b, "sell")) def test_swaps_of_different_mints_run_at_once_and_each_mints_in_order(circuit1): """Live money fast: one token's swap in flight never holds another token's back; one token's sale and buy-back still run one after the other.""" a, b = f"M{uuid.uuid4().hex[:10]}", f"M{uuid.uuid4().hex[:10]}" gate = threading.Event() fb = _ScriptBroker([_fill(-5, 5), _fill(-5, 5), _fill(5, -5)], gate=gate); w = ExecutionWorker(fb, threads=3) try: assert w.submit(_req(a, "sell")) and w.submit(_req(b, "sell")) and w.submit(_req(a, "buy", 5)) for _ in range(100): if len(fb.calls) >= 2: break threading.Event().wait(0.02) assert sorted((c["mint"], c["side"]) for c in fb.calls) == sorted([(a, "sell"), (b, "sell")]) # both sales in flight, a's buy behind its sale gate.set() got = [w.results.get(timeout=10) for _ in range(3)] assert [(c["mint"], c["side"]) for c in fb.calls if c["mint"] == a] == [(a, "sell"), (a, "buy")] # a's buy ran once its sale was done assert all(r.ok for r in got) and w.pending() == set() finally: gate.set(); w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = ANY(%s)", ([a, b],)) def test_a_wealth_mark_waits_for_every_swap_in_flight_and_never_holds_one_back(): from fly_trader.execution.worker import SharedLock busy = SharedLock() with busy.shared(): with busy.shared(): # two swaps in flight at once assert not busy.acquire(blocking=False) and not busy.acquire(timeout=0.05) assert busy.acquire(timeout=0.1) # none in flight: the mark reads alone started = threading.Event() def swap(): with busy.shared(): started.set() t = threading.Thread(target=swap); t.start() assert not started.wait(0.1) # a swap waits only while a mark holds it busy.release(); t.join(2) assert started.is_set() def test_the_worker_tells_its_mirror_once_an_order_is_booked(circuit1): """The live stats follow the live orders: the mirror's hook runs as soon as each order is booked, after its result.""" mint = f"M{uuid.uuid4().hex[:10]}"; seen = [] fb = _ScriptBroker([_fill(1_000_000, -200_000_000, 2e-7)]); w = ExecutionWorker(fb) w.on_done = lambda res: seen.append((res.request.side, res.ok, res.position_id is not None)) try: assert w.submit(_req(mint, "buy", 200_000_000)) w.results.get(timeout=10) for _ in range(100): if seen: break threading.Event().wait(0.05) assert seen == [("buy", True, True)] finally: w.stop() with transaction() as conn: conn.execute("DELETE FROM positions WHERE mint = %s", (mint,))