Skip to content

Commit eb11bd2

Browse files
nullchimpnullchimp
authored
Merge pull request #24 from nullchimp/session_management
Session management
2 parents 4027bd0 + a88244b commit eb11bd2

20 files changed

+584
-1257
lines changed

‎src/agent.py‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,12 @@
22
load_dotenv(override=True)
33

44
import asyncio
5-
5+
import os
66
from datetime import date
77

88
from core import chatutil, graceful_exit, pretty_print
99
from core.llm.chat import Chat
10+
from core.mcp.sessions_manager import MCPSessionManager
1011

1112
from tools import Tool
1213
from tools.github_search import GitHubKnowledgebase
@@ -33,6 +34,7 @@ def __init__(self, session_id: str = "cli-session"):
3334
tool.enable()
3435

3536
self.session_id = session_id
37+
self.mcp_initialized = False
3638

3739
self.history = []
3840
self._update_system_prompt()
@@ -116,6 +118,21 @@ def disable_tool(self, tool_name: str) -> bool:
116118
def get_tools(self) -> list:
117119
return self.chat.get_tools()
118120

121+
async def initialize_mcp_tools(self):
122+
if self.mcp_initialized:
123+
return
124+
125+
print("Initializing MCP tools...")
126+
127+
config_path = os.path.join(os.path.dirname(__file__), "..", "config", "mcp.json")
128+
session_manager = MCPSessionManager()
129+
await session_manager.discovery(config_path)
130+
for tool in session_manager.tools:
131+
self.add_tool(tool)
132+
133+
print("MCP tools initialized.")
134+
self.mcp_initialized = True
135+
119136
async def process_query(self, user_prompt: str) -> str:
120137
user_role = {"role": "user", "content": user_prompt}
121138

@@ -152,12 +169,17 @@ async def process_query(self, user_prompt: str) -> str:
152169
pretty_print("History", self.history)
153170
return result, tools_used
154171

155-
def get_agent_instance(session_id: str = None) -> Agent:
172+
async def get_agent_instance(session_id: str = None) -> Agent:
156173
if not session_id:
157174
raise ValueError("Session ID must be provided to get agent instance.")
158175

159176
if not session_id in _agent_sessions:
160-
_agent_sessions[session_id] = Agent(session_id)
177+
print(f"Creating new agent instance for session: {session_id}")
178+
agent: Agent = Agent(session_id)
179+
print(f"Agent instance created with session ID: {agent.session_id}")
180+
await agent.initialize_mcp_tools()
181+
182+
_agent_sessions[session_id] = agent
161183

162184
return _agent_sessions[session_id]
163185

@@ -172,13 +194,15 @@ def delete_agent_instance(session_id: str) -> bool:
172194
return False
173195

174196

175-
def add_tool(tool: Tool, session_id: str = "cli-session") -> None:
176-
get_agent_instance(session_id).add_tool(tool)
197+
async def add_tool(tool: Tool, session_id: str = "cli-session") -> None:
198+
agent = await get_agent_instance(session_id)
199+
agent.add_tool(tool)
177200

178201
@graceful_exit
179202
@chatutil("Agent")
180203
async def run_conversation(user_prompt, session_id: str = "cli-session") -> str:
181-
result = await get_agent_instance(session_id).process_query(user_prompt)
204+
agent = await get_agent_instance(session_id)
205+
result = await agent.process_query(user_prompt)
182206
pretty_print("Result", result)
183207
return result
184208

‎src/api/routes.py‎

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
from fastapi import APIRouter, Depends, HTTPException, Response
22
import uuid
3-
import os
43

54
from agent import get_agent_instance, Agent, delete_agent_instance
65
from api.auth import get_api_key
@@ -9,29 +8,23 @@
98
ToolToggleResponse, ToolInfo, DebugResponse, DebugRequest, NewSessionResponse
109
)
1110
from core.debug_capture import get_debug_capture_instance, get_all_debug_events, clear_all_debug_events, delete_debug_capture_instance
12-
from core.mcp.sessions_manager import MCPSessionManager
1311

1412
session_router = APIRouter(prefix="/api")
1513
api_router = APIRouter(prefix="/api/{session_id}", dependencies=[Depends(get_api_key)])
1614

17-
@session_router.post("/session/new", response_model=NewSessionResponse)
18-
async def new_session():
19-
session_id = str(uuid.uuid4())
20-
agent_instance = get_agent_instance(session_id)
21-
22-
get_debug_capture_instance(session_id)
23-
15+
@session_router.get("/session/{session_id}", response_model=NewSessionResponse)
16+
async def get_session(session_id: str):
2417
try:
25-
config_path = os.path.join(os.path.dirname(__file__), '..', '..', 'config', 'mcp.json')
26-
session_manager = MCPSessionManager()
27-
await session_manager.discovery(config_path)
28-
for tool in session_manager.tools:
29-
agent_instance.add_tool(tool)
18+
if session_id == "new":
19+
session_id = str(uuid.uuid4())
3020

31-
return NewSessionResponse(session_id=session_id, message="Session created successfully")
21+
await get_agent_instance(session_id)
22+
get_debug_capture_instance(session_id)
23+
return NewSessionResponse(session_id=session_id, message="Session is active")
3224
except Exception as e:
3325
raise HTTPException(status_code=500, detail=f"Error initializing agent: {str(e)}")
3426

27+
3528
@session_router.delete("/session/{session_id}")
3629
async def delete_session(session_id: str):
3730
if delete_agent_instance(session_id):

‎src/main.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ async def cli_main():
1818
session_manager = MCPSessionManager()
1919
await session_manager.discovery(path)
2020
for tool in session_manager.tools:
21-
agent.add_tool(tool)
21+
await agent.add_tool(tool)
2222

2323
await agent_task()
2424

‎src/ui/chat.ts‎

Lines changed: 98 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ class ChatApp {
6464
private debugEventsList: DebugEvent[] = [];
6565
private isCreatingSession: boolean = false;
6666
private isLoadingTools: boolean = false;
67+
private isSendingMessage: boolean = false;
68+
private isVerifyingSession: boolean = false;
6769

6870
private apiBaseUrl = 'http://localhost:5555/api';
6971

@@ -96,16 +98,25 @@ class ChatApp {
9698
this.loadChatHistory();
9799
this.setupEventListeners();
98100

101+
console.log('ChatApp initialized with', this.sessions.length, 'sessions loaded.');
102+
99103
// Only create a new session if no sessions exist
100104
if (this.sessions.length === 0) {
101105
await this.createNewSession();
102106
} else {
103107
// Load the most recent session
104108
this.currentSession = this.sessions[0];
105109

106-
// If the session doesn't have a backend sessionId, we'll create one when user sends a message
107-
if (!this.currentSession.sessionId) {
108-
console.log('Current session has no backend sessionId - will create when needed');
110+
console.log('Loading existing session:', this.currentSession.sessionId);
111+
// If the session has a backend sessionId, verify it exists
112+
if (this.currentSession.sessionId) {
113+
this.isVerifyingSession = true;
114+
this.updateSendButtonState();
115+
this.showSessionVerificationLoading();
116+
await this.verifyBackendSession(this.currentSession.sessionId);
117+
this.isVerifyingSession = false;
118+
this.updateSendButtonState();
119+
this.hideSessionVerificationLoading();
109120
}
110121

111122
await this.loadSession(this.currentSession.id);
@@ -169,7 +180,19 @@ class ChatApp {
169180
private updateSendButtonState(): void {
170181
const hasContent = this.messageInput.value.trim().length > 0;
171182
const hasActiveSession = this.currentSession !== null;
172-
this.sendBtn.disabled = !hasContent || !hasActiveSession;
183+
const isLoading = this.isLoadingTools || this.isSendingMessage || this.isCreatingSession || this.isVerifyingSession;
184+
185+
this.sendBtn.disabled = !hasContent || !hasActiveSession || isLoading;
186+
this.messageInput.disabled = isLoading;
187+
188+
// Add visual feedback for loading state
189+
if (isLoading) {
190+
this.messageInput.placeholder = 'Please wait...';
191+
this.messageInput.classList.add('loading');
192+
} else {
193+
this.messageInput.placeholder = 'Type your message...';
194+
this.messageInput.classList.remove('loading');
195+
}
173196
}
174197

175198
private updateNewChatButtonState(): void {
@@ -202,12 +225,15 @@ class ChatApp {
202225
const content = this.messageInput.value.trim();
203226
if (!content || !this.currentSession) return;
204227

205-
// If the current session doesn't have a backend session ID, create a new session
228+
this.isSendingMessage = true;
229+
this.updateSendButtonState();
230+
231+
// If the current session doesn't have a backend session ID, create one
206232
if (!this.currentSession.sessionId) {
207233
console.log('Creating backend session for existing frontend session...');
208234
try {
209235
const response = await fetch(`${this.apiBaseUrl}/session/new`, {
210-
method: 'POST',
236+
method: 'GET',
211237
headers: {
212238
'Content-Type': 'application/json'
213239
}
@@ -222,11 +248,15 @@ class ChatApp {
222248
await this.loadTools();
223249
} else {
224250
this.showError('Failed to create backend session. Please try again.');
251+
this.isSendingMessage = false;
252+
this.updateSendButtonState();
225253
return;
226254
}
227255
} catch (error) {
228256
console.error('Failed to create backend session:', error);
229257
this.showError('Failed to create backend session. Please try again.');
258+
this.isSendingMessage = false;
259+
this.updateSendButtonState();
230260
return;
231261
}
232262
}
@@ -273,6 +303,9 @@ class ChatApp {
273303
this.hideTypingIndicator();
274304
this.showError('Failed to get response. Please try again.');
275305
console.error('API Error:', error);
306+
} finally {
307+
this.isSendingMessage = false;
308+
this.updateSendButtonState();
276309
}
277310
}
278311

@@ -428,14 +461,15 @@ class ChatApp {
428461

429462
this.isCreatingSession = true;
430463
this.updateNewChatButtonState();
464+
this.updateSendButtonState();
431465

432466
// Show loading state immediately
433467
this.clearMessages();
434468

435469
try {
436470
// Create a new backend session
437471
const response = await fetch(`${this.apiBaseUrl}/session/new`, {
438-
method: 'POST',
472+
method: 'GET',
439473
headers: {
440474
'Content-Type': 'application/json'
441475
}
@@ -478,6 +512,7 @@ class ChatApp {
478512
} finally {
479513
this.isCreatingSession = false;
480514
this.updateNewChatButtonState();
515+
this.updateSendButtonState();
481516
// Clear the loading message and show the normal welcome message
482517
this.clearMessages();
483518
}
@@ -716,6 +751,7 @@ class ChatApp {
716751

717752
this.isLoadingTools = true;
718753
this.renderToolsLoading();
754+
this.updateSendButtonState();
719755

720756
const toolsUrl = `${this.apiBaseUrl}/${this.currentSession.sessionId}/tools`;
721757
console.log(`Loading tools from: ${toolsUrl}`);
@@ -741,6 +777,7 @@ class ChatApp {
741777
} finally {
742778
this.isLoadingTools = false;
743779
this.renderTools();
780+
this.updateSendButtonState();
744781
}
745782
}
746783

@@ -1258,9 +1295,9 @@ class ChatApp {
12581295
});
12591296
}
12601297

1298+
this.updateSendButtonState();
12611299
this.messageInput.disabled = true;
12621300
this.messageInput.placeholder = "Create a new chat to start messaging...";
1263-
this.updateSendButtonState();
12641301

12651302
// Clear tools and debug state for no active session
12661303
this.tools = [];
@@ -1279,6 +1316,59 @@ class ChatApp {
12791316
this.messageInput.placeholder = "Message AI Agent...";
12801317
this.updateSendButtonState();
12811318
}
1319+
1320+
private async verifyBackendSession(sessionId: string): Promise<void> {
1321+
try {
1322+
console.log(`Verifying backend session: ${sessionId}`);
1323+
const response = await fetch(`${this.apiBaseUrl}/session/${sessionId}`, {
1324+
method: 'GET',
1325+
headers: {
1326+
'Content-Type': 'application/json'
1327+
}
1328+
});
1329+
1330+
if (response.ok) {
1331+
const sessionData = await response.json();
1332+
console.log(`Backend session ${sessionId} verified and reinitialized`);
1333+
// Session exists and is reinitialized
1334+
return;
1335+
} else {
1336+
console.warn(`Backend session ${sessionId} not found, will create new session when needed`);
1337+
// Clear the sessionId since backend session doesn't exist
1338+
if (this.currentSession) {
1339+
this.currentSession.sessionId = undefined;
1340+
this.saveChatHistory();
1341+
}
1342+
}
1343+
} catch (error) {
1344+
console.error(`Failed to verify backend session ${sessionId}:`, error);
1345+
// Clear the sessionId since we can't verify
1346+
if (this.currentSession) {
1347+
this.currentSession.sessionId = undefined;
1348+
this.saveChatHistory();
1349+
}
1350+
}
1351+
}
1352+
1353+
private showSessionVerificationLoading(): void {
1354+
this.messagesContainer.innerHTML = `
1355+
<div class="welcome-message">
1356+
<div class="session-verification-loading">
1357+
<div class="loading-spinner"></div>
1358+
<h1>Verifying session...</h1>
1359+
<p>Checking if your previous session is still available.</p>
1360+
</div>
1361+
</div>
1362+
`;
1363+
}
1364+
1365+
private hideSessionVerificationLoading(): void {
1366+
// Clear the loading message - it will be replaced by the actual session content
1367+
const loadingEl = this.messagesContainer.querySelector('.session-verification-loading');
1368+
if (loadingEl) {
1369+
loadingEl.parentElement?.remove();
1370+
}
1371+
}
12821372
}
12831373

12841374
document.addEventListener('DOMContentLoaded', () => {

0 commit comments

Comments
 (0)