Update backend/main.py

backend/main.py

@@ -1,10 +1,4 @@
-
-# Fixed main.py
-main_py = '''import os
-import json
-import asyncio
-import subprocess
-import uuid
+import os, json, asyncio, subprocess, uuid
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks
 from fastapi.responses import FileResponse, StreamingResponse
@@ -26,13 +20,9 @@ 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:
@@ -43,20 +33,16 @@ 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: Dict[str, Any] = {}
-agents: Dict[str, CodingAgent] = {}
+indexers = {}
+agents = {}
 STATIC_DIR = "static"
-SANDBOX_ROOT = "/app/sandbox"
-
 
 # ---------- API Models ----------
 class ChatRequest(BaseModel):
@@ -65,15 +51,13 @@ 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 = "/app/sandbox"
+    project_path: str
     auto_approve: bool = False
     background: bool = False
 
-
 class AutocompleteRequest(BaseModel):
     prefix: str
     suffix: str = ""
@@ -82,101 +66,61 @@ class AutocompleteRequest(BaseModel):
     max_tokens: int = 64
     model: str = "groq-llama-3.1-8b-instant"
 
-
 class IndexRequest(BaseModel):
-    project_path: str = "/app/sandbox"
-
+    project_path: str
 
 class BugFixRequest(BaseModel):
     file_path: str
     error_output: str
-    project_path: str = "/app/sandbox"
-
+    project_path: str
 
 class ReviewRequest(BaseModel):
     pr_diff: str
-    project_path: str = "/app/sandbox"
-
+    project_path: str
 
 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:
-                    import git
-                    repo = git.Repo(safe_path)
+                    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 Exception:
+                    context += f"\n--- Git Diff ---\n{diff}\n"
+                except:
                     pass
             else:
-                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}"})
-
+                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:
-            async for token in route_chat(req.model, messages):
-                yield f"data: {json.dumps({'token': token})}\\n\\n"
-            yield "data: [DONE]\\n\\n"
+            async for token in route_chat(req.model, req.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):
-    safe_path = _validate_project_path(req.project_path)
-    agent = CodingAgent(project_path=safe_path, model=req.model)
+    agent = CodingAgent(project_path=req.project_path, model=req.model)
     changes = await agent.plan_task(req.prompt)
-    task_id = str(uuid.uuid4())
-    agents[task_id] = agent
-    return {"task_id": task_id, "changes": changes}
-
+    agents["default"] = agent
+    return {"changes": changes}
 
 @app.post("/api/agent/apply")
 async def agent_apply(req: ApplyRequest):
@@ -186,11 +130,9 @@ 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):
-    safe_path = _validate_project_path(req.project_path)
-    agent = CodingAgent(project_path=safe_path, model=req.model)
+    agent = CodingAgent(project_path=req.project_path, model=req.model)
     if req.background:
         task_id = str(uuid.uuid4())
         agents[task_id] = agent
@@ -199,226 +141,128 @@ 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(safe_path)
+    idx = CBI(req.project_path)
     idx.build_index()
-    indexers[safe_path] = idx
+    indexers[req.project_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):
-    safe_path = _validate_project_path(project_path)
-    if safe_path not in indexers:
+    if project_path not in indexers:
         CBI = get_codebase_index()
-        idx = CBI(safe_path)
+        idx = CBI(project_path)
         try:
             idx.build_index()
         except Exception as e:
             raise HTTPException(503, f"Indexing failed: {str(e)}")
-        indexers[safe_path] = idx
-
-    idx = indexers[safe_path]
+        indexers[project_path] = idx
+    idx = indexers[project_path]
     results = idx.search(query, top_k)
     return {"results": results}
 
-
 @app.post("/api/bugfix")
 async def bugfix(req: BugFixRequest):
-    safe_path = _validate_project_path(req.project_path)
-    agent = CodingAgent(
-        project_path=safe_path,
-        model="openrouter-anthropic/claude-3-sonnet"
-    )
+    agent = CodingAgent(project_path=req.project_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):
-    safe_path = _validate_project_path(req.project_path)
-    agent = CodingAgent(
-        project_path=safe_path,
-        model="openrouter-anthropic/claude-3-sonnet"
-    )
+    agent = CodingAgent(project_path=req.project_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(req: GitRequest):
-    safe_path = _validate_project_path(req.project_path)
+async def git_operation(action: str, project_path: str, args: Optional[dict] = None):
+    import git
     try:
-        import git
-        repo = git.Repo(safe_path)
-        args = req.args or {}
-
-        if req.action == "commit":
+        repo = git.Repo(project_path)
+        if action == "commit":
             repo.git.add(A=True)
             repo.git.commit(m=args.get("message", "auto commit"))
-            return {"status": "committed"}
-
-        elif req.action == "log":
+        elif action == "log":
             logs = list(repo.iter_commits(max_count=10))
-            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":
+            return {"log": [{"message": c.message, "hash": c.hexsha} for c in logs]}
+        elif action == "branch":
             branches = [b.name for b in repo.branches]
             return {"branches": branches}
-
-        elif req.action == "diff":
-            diff_output = ""
-            if repo.head.is_valid():
-                diff_output = repo.git.diff(repo.head.commit.tree)
+        elif action == "diff":
+            diff_output = repo.git.diff(repo.head.commit.tree) if repo.head.is_valid() else ""
             return {"diff": diff_output}
-
-        elif req.action == "status":
-            return {"status": repo.git.status()}
-
         else:
-            raise HTTPException(400, f"Unknown git action: {req.action}")
-
+            raise HTTPException(400, "Unknown action")
     except Exception as e:
-        raise HTTPException(400, detail=str(e))
-
+        raise HTTPException(400, 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=SANDBOX_ROOT,
-                )
-                # Stream stdout and stderr concurrently
-                await asyncio.gather(
-                    read_stream(proc.stdout, websocket),
-                    read_stream(proc.stderr, websocket),
+                    cwd="/app/sandbox"
                 )
-                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:
+                stdout, stderr = await proc.communicate()
+                output = stdout.decode() + stderr.decode()
+                await websocket.send_text(output)
+            elif cmd == "ctrl-c" and proc:
                 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.get("/api/agent/task/{task_id}")
+@app.post("/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"):
-    safe_path = _validate_project_path(project_path)
-
-    def build_tree(path: str):
+    def build_tree(path):
         tree = []
-        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
+        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)
         return tree
-
     try:
-        tree = build_tree(safe_path)
+        tree = build_tree(project_path)
         return tree
     except Exception as e:
         raise HTTPException(400, str(e))
 
-
-# ---------- 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")
-
-
+# ---------- Serve Frontend (catch-all, must be last) ----------
 @app.get("/{full_path:path}")
 async def serve_frontend(full_path: str):
-    if full_path.startswith("api/") or full_path.startswith("ws/"):
+    if full_path.startswith("api"):
         raise HTTPException(404)
     file_path = os.path.join(STATIC_DIR, full_path)
     if os.path.isfile(file_path):
@@ -428,12 +272,6 @@ 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)
-'''
-
-with open('/mnt/agents/output/main.py', 'w') as f:
-    f.write(main_py)
-print("main.py written successfully")
+    uvicorn.run(app, host="0.0.0.0", port=7860)