Update backend/main.py

backend/main.py

@@ -1,5 +1,6 @@
 import os, json, asyncio, subprocess, uuid
-from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks, Request
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks
 from fastapi.responses import FileResponse, StreamingResponse
 from pydantic import BaseModel
 from typing import List, Optional, Dict, Any
@@ -22,10 +23,25 @@ def get_codebase_index():
             raise HTTPException(status_code=503, detail="Codebase indexing unavailable – rebuild with updated sentence-transformers")
     return CodebaseIndex
 
-app = FastAPI(title="Codeki API")
+# Auto-index sandbox on startup
+async def startup_index():
+    try:
+        from codebase_index import CodebaseIndex
+        idx = CodebaseIndex("/app/sandbox")
+        idx.build_index()
+        indexers["/app/sandbox"] = idx
+    except Exception as e:
+        print(f"Startup indexing failed: {e}")
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    await startup_index()
+    yield
 
-indexers = {}     # project_path -> CodebaseIndex instance
-agents = {}       # task_id -> CodingAgent instance
+app = FastAPI(title="Codeki API", lifespan=lifespan)
+
+indexers = {}
+agents = {}
 STATIC_DIR = "static"
 
 # ---------- API Models ----------
@@ -62,6 +78,10 @@ class ReviewRequest(BaseModel):
     pr_diff: str
     project_path: str
 
+class ApplyRequest(BaseModel):
+    task_id: str
+    approved_indices: List[int]
+
 # ---------- API Routes ----------
 @app.get("/api/models")
 async def get_models():
@@ -72,9 +92,18 @@ async def chat(req: ChatRequest):
     context = ""
     if req.context_files:
         for f in req.context_files:
-            p = Path(req.project_path) / f
-            if p.exists():
-                context += f"\n--- {f} ---\n{p.read_text()}\n"
+            if f == "__git__":
+                import git
+                try:
+                    repo = git.Repo(req.project_path)
+                    diff = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
+                    context += f"\n--- Git Diff ---\n{diff}\n"
+                except:
+                    pass
+            else:
+                p = Path(req.project_path) / f
+                if p.exists():
+                    context += f"\n--- {f} ---\n{p.read_text()}\n"
         req.messages.insert(0, {"role": "system", "content": f"Relevant codebase context:\n{context}"})
     async def event_stream():
         try:
@@ -86,6 +115,21 @@ async def chat(req: ChatRequest):
             yield "data: [DONE]\n\n"
     return StreamingResponse(event_stream(), media_type="text/event-stream")
 
+@app.post("/api/agent/plan")
+async def agent_plan(req: AgentTaskRequest):
+    agent = CodingAgent(project_path=req.project_path, model=req.model)
+    changes = await agent.plan_task(req.prompt)
+    agents["default"] = agent
+    return {"changes": changes}
+
+@app.post("/api/agent/apply")
+async def agent_apply(req: ApplyRequest):
+    agent = agents.get(req.task_id)
+    if not agent:
+        raise HTTPException(404, "Agent not found")
+    await agent.apply_changes(req.approved_indices)
+    return {"status": agent.status, "logs": agent.logs}
+
 @app.post("/api/agent/run")
 async def run_agent(req: AgentTaskRequest, background_tasks: BackgroundTasks):
     agent = CodingAgent(project_path=req.project_path, model=req.model)
@@ -117,7 +161,6 @@ async def index_project(req: IndexRequest):
 @app.post("/api/search")
 async def search_codebase(query: str, project_path: str = "/app/sandbox", top_k: int = 5):
     if project_path not in indexers:
-        # Auto-build index if not present
         CBI = get_codebase_index()
         idx = CBI(project_path)
         try:
@@ -156,7 +199,6 @@ async def git_operation(action: str, project_path: str, args: Optional[dict] = N
             branches = [b.name for b in repo.branches]
             return {"branches": branches}
         elif action == "diff":
-            # Return the diff of unstaged changes, or HEAD if no changes
             diff_output = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
             return {"diff": diff_output}
         else:
@@ -200,18 +242,13 @@ async def get_agent_status(task_id: str):
 
 @app.get("/api/files")
 async def list_files(project_path: str = "/app/sandbox"):
-    """Return a tree of the project folder"""
     def build_tree(path):
         tree = []
         for item in sorted(os.listdir(path)):
             if item.startswith('.') or item == '__pycache__':
                 continue
             full = os.path.join(path, item)
-            node = {
-                "name": item,
-                "path": full.replace(project_path, "").lstrip("/"),
-                "type": "folder" if os.path.isdir(full) else "file"
-            }
+            node = {"name": item, "path": full.replace(project_path, "").lstrip("/"), "type": "folder" if os.path.isdir(full) else "file"}
             if os.path.isdir(full):
                 node["children"] = build_tree(full)
             tree.append(node)