118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
import time
|
|
from pathlib import Path
|
|
from typing import Optional, Callable, Dict, Any, List
|
|
import httpx
|
|
|
|
from agent.config import AgentConfig, load_config
|
|
from agent.chunker import FileChunker
|
|
from agent.state_db import state_db
|
|
|
|
class ChunkUploader:
|
|
"""HTTP/HTTPS Chunk transfer client with automatic resumption and retry backoff."""
|
|
|
|
def __init__(self, config: Optional[AgentConfig] = None):
|
|
self.config = config or load_config()
|
|
|
|
def _get_headers(self) -> Dict[str, str]:
|
|
if not self.config.device_id or not self.config.device_token:
|
|
raise ValueError("Agent is not registered. Run 'agent_cli.py register' first.")
|
|
return {
|
|
"X-Device-Id": self.config.device_id,
|
|
"X-Device-Token": self.config.device_token
|
|
}
|
|
|
|
def upload_file(
|
|
self,
|
|
filepath: Path,
|
|
job_id: Optional[int] = None,
|
|
progress_callback: Optional[Callable[[int, int, float], None]] = None,
|
|
max_retries: int = 5
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Uploads a file chunk by chunk. If interrupted, querying the server will return
|
|
already received chunks, allowing immediate resumption without re-uploading completed parts.
|
|
"""
|
|
filepath = Path(filepath).resolve()
|
|
chunker = FileChunker(filepath=filepath, chunk_size=self.config.chunk_size)
|
|
full_sha256 = chunker.full_sha256
|
|
filename = filepath.name
|
|
|
|
headers = self._get_headers()
|
|
base_url = self.config.server_url.rstrip("/")
|
|
|
|
with httpx.Client(base_url=base_url, headers=headers, timeout=60.0) as client:
|
|
# 1. Initialize or resume upload session
|
|
init_payload = {
|
|
"filename": filename,
|
|
"file_size": chunker.file_size,
|
|
"sha256": full_sha256,
|
|
"chunk_size": chunker.chunk_size,
|
|
"job_id": job_id
|
|
}
|
|
|
|
resp = client.post("/api/upload/session", json=init_payload)
|
|
resp.raise_for_status()
|
|
session_data = resp.json()
|
|
|
|
session_code = session_data["session_code"]
|
|
total_chunks = session_data["total_chunks"]
|
|
received_chunks = set(session_data.get("received_chunks", []))
|
|
|
|
state_db.save_session(
|
|
filepath=str(filepath),
|
|
sha256=full_sha256,
|
|
session_code=session_code,
|
|
total_chunks=total_chunks,
|
|
chunk_size=chunker.chunk_size,
|
|
status="UPLOADING"
|
|
)
|
|
|
|
# 2. Upload missing chunks
|
|
missing_chunks = [i for i in range(total_chunks) if i not in received_chunks]
|
|
|
|
for idx in missing_chunks:
|
|
chunk_bytes, chunk_hash = chunker.get_chunk(idx)
|
|
|
|
# Retry loop with exponential backoff for network resilience
|
|
success = False
|
|
attempt = 0
|
|
while not success and attempt < max_retries:
|
|
try:
|
|
chunk_headers = {
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Chunk-Index": str(idx),
|
|
"X-Chunk-SHA256": chunk_hash
|
|
}
|
|
c_resp = client.post(
|
|
f"/api/upload/{session_code}/chunk?chunk_index={idx}&chunk_sha256={chunk_hash}",
|
|
content=chunk_bytes,
|
|
headers=chunk_headers
|
|
)
|
|
c_resp.raise_for_status()
|
|
success = True
|
|
except Exception as ex:
|
|
attempt += 1
|
|
if attempt >= max_retries:
|
|
raise ConnectionError(f"Failed to upload chunk {idx} after {max_retries} attempts: {str(ex)}")
|
|
time.sleep(2 ** attempt)
|
|
|
|
received_chunks.add(idx)
|
|
if progress_callback:
|
|
pct = round((len(received_chunks) / total_chunks) * 100, 2)
|
|
progress_callback(len(received_chunks), total_chunks, pct)
|
|
|
|
# 3. Complete and verify full file assembly on server
|
|
complete_resp = client.post(f"/api/upload/{session_code}/complete")
|
|
complete_resp.raise_for_status()
|
|
result = complete_resp.json()
|
|
|
|
# 4. Update local state database
|
|
state_db.mark_session_completed(
|
|
filepath=str(filepath),
|
|
sha256=full_sha256,
|
|
file_size=chunker.file_size,
|
|
job_id=job_id
|
|
)
|
|
|
|
return result
|