Yield Agent

from agent_sdk import Agent, AgentContext
from agent_sdk.agent_context import CurrentPosition
from eth_abi import encode

# Base addresses
USDC = "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913"
AAVE_V3_POOL = "0xA238Dd80C259a72e81d7e4664a9801593F98d1c5"

# Function selectors
APPROVE_SELECTOR = "095ea7b3"
SUPPLY_SELECTOR = "617ba037"
WITHDRAW_SELECTOR = "69328dec"
MAX_UINT256 = 2**256 - 1

def build_approve_data(spender: str, amount: int) -> str:
    params = encode(["address", "uint256"], [spender, amount])
    return "0x" + APPROVE_SELECTOR + params.hex()

def build_supply_data(asset: str, amount: int, on_behalf_of: str) -> str:
    params = encode(
        ["address", "uint256", "address", "uint16"], [asset, amount, on_behalf_of, 0]
    )
    return "0x" + SUPPLY_SELECTOR + params.hex()

def build_withdraw_data(asset: str, amount: int, to: str) -> str:
    params = encode(["address", "uint256", "address"], [asset, amount, to])
    return "0x" + WITHDRAW_SELECTOR + params.hex()

def run(agent: AgentContext) -> None:
    # Find USDC balance
    usdc = next(
        (p for p in agent.currentPositions if p.assetAddress.lower() == USDC.lower()),
        None,
    )
    if not usdc or int(usdc.currentQty) == 0:
        agent.log("No USDC to deposit")
        return

    amount = int(usdc.currentQty)
    agent.log(f"Depositing {usdc.currentQty} USDC (raw) into Aave V3")

    # Step 1: Approve Aave Pool to spend USDC (skip if already approved)
    approved_check = agent.memory.get("aaveApproved")
    already_approved = approved_check.success and approved_check.data and approved_check.data.value == "true"

    if not already_approved:
        approve_result = agent.sign_and_send({
            "network": "ethereum:8453",
            "request": {
                "to_address": USDC,
                "data": build_approve_data(AAVE_V3_POOL, MAX_UINT256),
                "value": "0",
            },
            "message": "Approve USDC for Aave V3",
        })

        if not approve_result.success:
            agent.log(f"Approve failed: {approve_result.error}", error=True)
            return
        agent.memory.set("aaveApproved", "true")
        agent.log("USDC approved for Aave V3")

    # Step 2: Supply USDC to Aave V3
    supply_result = agent.sign_and_send({
        "network": "ethereum:8453",
        "request": {
            "to_address": AAVE_V3_POOL,
            "data": build_supply_data(USDC, amount, agent.sessionWalletAddress),
            "value": "0",
        },
        "message": "Supply USDC to Aave V3",
    })

    if supply_result.success:
        agent.log("USDC deposited into Aave V3!")
    else:
        agent.log(f"Supply failed: {supply_result.error}", error=True)

def unwind(agent: AgentContext, positions: list[CurrentPosition]) -> None:
    # Withdraw all USDC from Aave V3
    result = agent.sign_and_send({
        "network": "ethereum:8453",
        "request": {
            "to_address": AAVE_V3_POOL,
            "data": build_withdraw_data(USDC, MAX_UINT256, agent.sessionWalletAddress),
            "value": "0",
        },
        "message": "Withdraw all USDC from Aave V3",
    })

    if result.success:
        agent.log("Withdrawn all USDC from Aave V3")
    else:
        agent.log(f"Withdraw failed: {result.error}", error=True)

agent = Agent(run_function=run, unwind_function=unwind)

handler = agent.get_handler()

if __name__ == "__main__":
    agent.run()