diff --git a/datasource/flask_api_source.py b/datasource/flask_api_source.py
index 87eb40c..a2368b2 100644
--- a/datasource/flask_api_source.py
+++ b/datasource/flask_api_source.py
@@ -146,8 +146,11 @@ class FlaskAPIDataSource:
                     df['date'] = pd.to_datetime(df['date'])
                     df = df.set_index('date')
                 
-                # 确保列名标准化
-                df = df[['open', 'high', 'low', 'close', 'volume']]
+                # 确保列名标准化（保留 code 列如果存在）
+                standard_cols = ['open', 'high', 'low', 'close', 'volume']
+                if 'code' in df.columns:
+                    standard_cols = ['code'] + standard_cols
+                df = df[standard_cols]
                 
                 # 使用 API 返回的实际数据范围（而非请求参数）
                 actual_start = data.get('date_range', {}).get('start', start_date)