Update backend/main.py

backend/main.py

@@ -1,4 +1,10 @@
-import os, json, asyncio, subprocess, uuid
+
+# Fixed main.py
+main_py = '''import os
+import json
+import asyncio
+import subprocess
+import uuid
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks
 from fastapi.responses import FileResponse, StreamingResponse
@@ -20,9 +26,13 @@ def get_codebase_index():
             from codebase_index import CodebaseIndex as CBI
             CodebaseIndex = CBI
         except ImportError:
-            raise HTTPException(status_code=503, detail="Codebase indexing unavailable – rebuild with updated sentence-transformers")
+            raise HTTPException(
+                status_code=503,
+                detail="Codebase indexing unavailable – rebuild with updated sentence-transformers"
+            )
     return CodebaseIndex
 
+
 # Auto-index sandbox on startup
 async def startup_index():
     try:
@@ -33,16 +43,20 @@ async def startup_index():
     except Exception as e:
         print(f"Startup indexing failed: {e}")
 
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     await startup_index()
     yield
 
+
 app = FastAPI(title="Codeki API", lifespan=lifespan)
 
-indexers = {}
-agents = {}
+indexers: Dict[str, Any] = {}
+agents: Dict[str, CodingAgent] = {}
 STATIC_DIR = "static"
+SANDBOX_ROOT = "/app/sandbox"
+
 
 # ---------- API Models ----------
 class ChatRequest(BaseModel):
@@ -51,13 +65,15 @@ class ChatRequest(BaseModel):
     context_files: Optional[List[str]] = None
     project_path: str = "/app/sandbox"
 
+
 class AgentTaskRequest(BaseModel):
     prompt: str
     model: str = "openrouter-anthropic/claude-3-sonnet"
-    project_path: str
+    project_path: str = "/app/sandbox"
     auto_approve: bool = False
     background: bool = False
 
+
 class AutocompleteRequest(BaseModel):
     prefix: str
     suffix: str = ""
@@ -66,61 +82,101 @@ class AutocompleteRequest(BaseModel):
     max_tokens: int = 64
     model: str = "groq-llama-3.1-8b-instant"
 
+
 class IndexRequest(BaseModel):
-    project_path: str
+    project_path: str = "/app/sandbox"
+
 
 class BugFixRequest(BaseModel):
     file_path: str
     error_output: str
-    project_path: str
+    project_path: str = "/app/sandbox"
+
 
 class ReviewRequest(BaseModel):
     pr_diff: str
-    project_path: str
+    project_path: str = "/app/sandbox"
+
 
 class ApplyRequest(BaseModel):
     task_id: str
     approved_indices: List[int]
 
+
+class GitRequest(BaseModel):
+    action: str
+    project_path: str = "/app/sandbox"
+    args: Optional[dict] = None
+
+
+def _validate_project_path(project_path: str) -> str:
+    """Validate and resolve project path to prevent directory traversal."""
+    real_path = os.path.realpath(project_path)
+    real_sandbox = os.path.realpath(SANDBOX_ROOT)
+    if not real_path.startswith(real_sandbox):
+        raise HTTPException(status_code=403, detail="Access denied: path outside sandbox")
+    if not os.path.exists(real_path):
+        raise HTTPException(status_code=404, detail="Project path not found")
+    return real_path
+
+
 # ---------- API Routes ----------
 @app.get("/api/models")
 async def get_models():
     return get_available_models()
 
+
 @app.post("/api/chat")
 async def chat(req: ChatRequest):
     context = ""
+    safe_path = _validate_project_path(req.project_path)
+
     if req.context_files:
         for f in req.context_files:
             if f == "__git__":
-                import git
                 try:
-                    repo = git.Repo(req.project_path)
+                    import git
+                    repo = git.Repo(safe_path)
                     diff = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
-                    context += f"\n--- Git Diff ---\n{diff}\n"
-                except:
+                    context += f"\\n--- Git Diff ---\\n{diff}\\n"
+                except Exception:
                     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}"})
+                p = Path(safe_path) / f
+                try:
+                    real_p = os.path.realpath(p)
+                    if not real_p.startswith(os.path.realpath(SANDBOX_ROOT)):
+                        continue
+                    if p.exists() and p.is_file():
+                        context += f"\\n--- {f} ---\\n{p.read_text(encoding='utf-8')}\\n"
+                except Exception:
+                    pass
+
+    messages = req.messages.copy()
+    if context:
+        messages.insert(0, {"role": "system", "content": f"Relevant codebase context:\\n{context}"})
+
     async def event_stream():
         try:
-            async for token in route_chat(req.model, req.messages):
-                yield f"data: {json.dumps({'token': token})}\n\n"
-            yield "data: [DONE]\n\n"
+            async for token in route_chat(req.model, messages):
+                yield f"data: {json.dumps({'token': token})}\\n\\n"
+            yield "data: [DONE]\\n\\n"
         except Exception as e:
-            yield f"data: {json.dumps({'token': f'ERROR: {str(e)}'})}\n\n"
-            yield "data: [DONE]\n\n"
+            yield f"data: {json.dumps({'token': f'ERROR: {str(e)}'})}\\n\\n"
+            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)
+    safe_path = _validate_project_path(req.project_path)
+    agent = CodingAgent(project_path=safe_path, model=req.model)
     changes = await agent.plan_task(req.prompt)
-    agents["default"] = agent
-    return {"changes": changes}
+    task_id = str(uuid.uuid4())
+    agents[task_id] = agent
+    return {"task_id": task_id, "changes": changes}
+
 
 @app.post("/api/agent/apply")
 async def agent_apply(req: ApplyRequest):
@@ -130,9 +186,11 @@ async def agent_apply(req: ApplyRequest):
     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)
+    safe_path = _validate_project_path(req.project_path)
+    agent = CodingAgent(project_path=safe_path, model=req.model)
     if req.background:
         task_id = str(uuid.uuid4())
         agents[task_id] = agent
@@ -141,128 +199,226 @@ async def run_agent(req: AgentTaskRequest, background_tasks: BackgroundTasks):
     else:
         async def event_stream():
             async for event in agent.execute_task_stream(req.prompt, auto_approve=req.auto_approve):
-                yield f"data: {json.dumps(event)}\n\n"
-            yield "data: [DONE]\n\n"
+                yield f"data: {json.dumps(event)}\\n\\n"
+            yield "data: [DONE]\\n\\n"
         return StreamingResponse(event_stream(), media_type="text/event-stream")
 
+
 @app.post("/api/autocomplete")
 async def autocomplete(req: AutocompleteRequest):
     suggestions = await generate_autocomplete(req)
     return {"completions": suggestions}
 
+
 @app.post("/api/index")
 async def index_project(req: IndexRequest):
+    safe_path = _validate_project_path(req.project_path)
     CBI = get_codebase_index()
-    idx = CBI(req.project_path)
+    idx = CBI(safe_path)
     idx.build_index()
-    indexers[req.project_path] = idx
+    indexers[safe_path] = idx
     return {"status": "indexed", "files": len(idx.file_list)}
 
+
 @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:
+    safe_path = _validate_project_path(project_path)
+    if safe_path not in indexers:
         CBI = get_codebase_index()
-        idx = CBI(project_path)
+        idx = CBI(safe_path)
         try:
             idx.build_index()
         except Exception as e:
             raise HTTPException(503, f"Indexing failed: {str(e)}")
-        indexers[project_path] = idx
-    idx = indexers[project_path]
+        indexers[safe_path] = idx
+
+    idx = indexers[safe_path]
     results = idx.search(query, top_k)
     return {"results": results}
 
+
 @app.post("/api/bugfix")
 async def bugfix(req: BugFixRequest):
-    agent = CodingAgent(project_path=req.project_path, model="openrouter-anthropic/claude-3-sonnet")
+    safe_path = _validate_project_path(req.project_path)
+    agent = CodingAgent(
+        project_path=safe_path,
+        model="openrouter-anthropic/claude-3-sonnet"
+    )
     result = await agent.fix_error(req.file_path, req.error_output)
     return result
 
+
 @app.post("/api/review")
 async def review_code(req: ReviewRequest):
-    agent = CodingAgent(project_path=req.project_path, model="openrouter-anthropic/claude-3-sonnet")
+    safe_path = _validate_project_path(req.project_path)
+    agent = CodingAgent(
+        project_path=safe_path,
+        model="openrouter-anthropic/claude-3-sonnet"
+    )
     review = await agent.review_pr(req.pr_diff)
     return {"review": review}
 
+
 @app.post("/api/git")
-async def git_operation(action: str, project_path: str, args: Optional[dict] = None):
-    import git
+async def git_operation(req: GitRequest):
+    safe_path = _validate_project_path(req.project_path)
     try:
-        repo = git.Repo(project_path)
-        if action == "commit":
+        import git
+        repo = git.Repo(safe_path)
+        args = req.args or {}
+
+        if req.action == "commit":
             repo.git.add(A=True)
             repo.git.commit(m=args.get("message", "auto commit"))
-        elif action == "log":
+            return {"status": "committed"}
+
+        elif req.action == "log":
             logs = list(repo.iter_commits(max_count=10))
-            return {"log": [{"message": c.message, "hash": c.hexsha} for c in logs]}
-        elif action == "branch":
+            return {
+                "log": [
+                    {
+                        "message": c.message.strip(),
+                        "hash": c.hexsha,
+                        "author": str(c.author),
+                        "date": str(c.committed_datetime),
+                    }
+                    for c in logs
+                ]
+            }
+
+        elif req.action == "branch":
             branches = [b.name for b in repo.branches]
             return {"branches": branches}
-        elif action == "diff":
-            diff_output = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
+
+        elif req.action == "diff":
+            diff_output = ""
+            if repo.head.is_valid():
+                diff_output = repo.git.diff(repo.head.commit.tree)
             return {"diff": diff_output}
+
+        elif req.action == "status":
+            return {"status": repo.git.status()}
+
         else:
-            raise HTTPException(400, "Unknown action")
+            raise HTTPException(400, f"Unknown git action: {req.action}")
+
     except Exception as e:
-        raise HTTPException(400, str(e))
+        raise HTTPException(400, detail=str(e))
+
 
 @app.websocket("/ws/terminal")
 async def terminal_ws(websocket: WebSocket):
     await websocket.accept()
     proc = None
+
+    async def read_stream(stream, websocket: WebSocket):
+        try:
+            while True:
+                line = await stream.readline()
+                if not line:
+                    break
+                await websocket.send_text(line.decode(errors='replace'))
+        except Exception:
+            pass
+
     try:
         while True:
             cmd = await websocket.receive_text()
+
             if cmd.startswith("exec:"):
                 command = cmd[5:]
                 proc = await asyncio.create_subprocess_shell(
                     command,
                     stdout=asyncio.subprocess.PIPE,
                     stderr=asyncio.subprocess.PIPE,
-                    cwd="/app/sandbox"
+                    cwd=SANDBOX_ROOT,
+                )
+                # Stream stdout and stderr concurrently
+                await asyncio.gather(
+                    read_stream(proc.stdout, websocket),
+                    read_stream(proc.stderr, websocket),
                 )
-                stdout, stderr = await proc.communicate()
-                output = stdout.decode() + stderr.decode()
-                await websocket.send_text(output)
-            elif cmd == "ctrl-c" and proc:
+                await proc.wait()
+                await websocket.send_text(f"\\n[exit code: {proc.returncode}]\\n")
+                proc = None
+
+            elif cmd == "ctrl-c" and proc and proc.returncode is None:
                 proc.terminate()
+                try:
+                    await asyncio.wait_for(proc.wait(), timeout=2.0)
+                except asyncio.TimeoutError:
+                    proc.kill()
+                proc = None
+
     except WebSocketDisconnect:
         pass
+    finally:
+        if proc and proc.returncode is None:
+            proc.kill()
+            try:
+                await asyncio.wait_for(proc.wait(), timeout=1.0)
+            except asyncio.TimeoutError:
+                pass
+
 
 @app.get("/api/agent/tasks")
 async def list_agent_tasks():
     return {"tasks": list(agents.keys())}
 
-@app.post("/api/agent/task/{task_id}")
+
+@app.get("/api/agent/task/{task_id}")
 async def get_agent_status(task_id: str):
     agent = agents.get(task_id)
     if not agent:
         raise HTTPException(404, "Task not found")
     return {"status": agent.status, "logs": agent.logs}
 
+
 @app.get("/api/files")
 async def list_files(project_path: str = "/app/sandbox"):
-    def build_tree(path):
+    safe_path = _validate_project_path(project_path)
+
+    def build_tree(path: str):
         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"}
-            if os.path.isdir(full):
-                node["children"] = build_tree(full)
-            tree.append(node)
+        try:
+            for item in sorted(os.listdir(path)):
+                if item.startswith('.') or item == '__pycache__':
+                    continue
+                full = os.path.join(path, item)
+                real_full = os.path.realpath(full)
+                if not real_full.startswith(os.path.realpath(SANDBOX_ROOT)):
+                    continue
+                node = {
+                    "name": item,
+                    "path": full.replace(safe_path, "").lstrip("/"),
+                    "type": "folder" if os.path.isdir(full) else "file",
+                }
+                if os.path.isdir(full):
+                    node["children"] = build_tree(full)
+                tree.append(node)
+        except PermissionError:
+            pass
         return tree
+
     try:
-        tree = build_tree(project_path)
+        tree = build_tree(safe_path)
         return tree
     except Exception as e:
         raise HTTPException(400, str(e))
 
-# ---------- Serve Frontend (catch-all, must be last) ----------
+
+# ---------- Serve Frontend ----------
+@app.get("/")
+async def serve_index():
+    index_path = os.path.join(STATIC_DIR, "index.html")
+    if os.path.isfile(index_path):
+        return FileResponse(index_path)
+    raise HTTPException(404, detail="Frontend not built")
+
+
 @app.get("/{full_path:path}")
 async def serve_frontend(full_path: str):
-    if full_path.startswith("api"):
+    if full_path.startswith("api/") or full_path.startswith("ws/"):
         raise HTTPException(404)
     file_path = os.path.join(STATIC_DIR, full_path)
     if os.path.isfile(file_path):
@@ -272,6 +428,12 @@ async def serve_frontend(full_path: str):
         return FileResponse(index_path)
     raise HTTPException(404, detail="Not found")
 
+
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+'''
+
+with open('/mnt/agents/output/main.py', 'w') as f:
+    f.write(main_py)
+print("main.py written successfully")