tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

feat: sessions api, view & resume prev sessions (#1453)

Browse Source 
* Centralize session files to goose::session module
* Write session metadata and messages in jsonl
* Refactor CLI build_session to use goose::session functions
* Track session's token usage by adding optional session_id in agent.reply(...)
* NOTE: Only sessions saved through the updates goose::session functions will show up in GUI

Co-authored-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
Salman Mohammed
2025-03-03 11:49:15 -05:00
committed by GitHub
parent 68b8c5d19d
commit 9ae9045584
25 changed files with 1413 additions and 257 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/mod.rs
+3 -1
View File
@@ -5,6 +5,7 @@ pub mod configs;
pub mod extension;
pub mod health;
pub mod reply;
pub mod session;
use axum::Router;
@@ -16,5 +17,6 @@ pub fn configure(state: crate::state::AppState) -> Router {
.merge(agent::routes(state.clone()))
.merge(extension::routes(state.clone()))
.merge(configs::routes(state.clone()))
.merge(config_management::routes(state))
.merge(config_management::routes(state.clone()))
.merge(session::routes(state))
}
crates/goose-server/src/routes/reply.rs
+73 -3
View File
@@ -9,6 +9,7 @@ use axum::{
use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::message::{Message, MessageContent};
use goose::session;
use mcp_core::role::Role;
use serde::{Deserialize, Serialize};
@@ -27,6 +28,7 @@ use tokio_stream::wrappers::ReceiverStream;
#[derive(Debug, Deserialize)]
struct ChatRequest {
messages: Vec<Message>,
session_id: Option<String>,
}
// Custom SSE response type for streaming messages
@@ -109,6 +111,11 @@ async fn handler(
// Get messages directly from the request
let messages = request.messages;
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
// Get a lock on the shared agent
let agent = state.agent.clone();
@@ -136,7 +143,16 @@ async fn handler(
}
};
let mut stream = match agent.reply(&messages).await {
// Get the provider first, before starting the reply stream
let provider = agent.provider().await;
let mut stream = match agent
.reply(
&messages,
Some(session::Identifier::Name(session_id.clone())),
)
.await
{
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
@@ -158,11 +174,16 @@ async fn handler(
}
};
// Collect all messages for storage
let mut all_messages = messages.clone();
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
loop {
tokio::select! {
response = timeout(Duration::from_millis(500), stream.next()) => {
match response {
Ok(Some(Ok(message))) => {
all_messages.push(message.clone());
if let Err(e) = stream_event(MessageEvent::Message { message }, &tx).await {
tracing::error!("Error sending message through channel: {}", e);
let _ = stream_event(
@@ -173,6 +194,16 @@ async fn handler(
).await;
break;
}
// Store messages and generate description in background
let session_path = session_path.clone();
let messages = all_messages.clone();
let provider = provider.clone();
tokio::spawn(async move {
if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
tracing::error!("Failed to store session history: {:?}", e);
}
});
}
Ok(Some(Err(e))) => {
tracing::error!("Error processing message: {}", e);
@@ -214,6 +245,7 @@ async fn handler(
#[derive(Debug, Deserialize, Serialize)]
struct AskRequest {
prompt: String,
session_id: Option<String>,
}
#[derive(Debug, Serialize)]
@@ -237,16 +269,30 @@ async fn ask_handler(
return Err(StatusCode::UNAUTHORIZED);
}
// Generate a new session ID if not provided in the request
let session_id = request
.session_id
.unwrap_or_else(session::generate_session_id);
let agent = state.agent.clone();
let agent = agent.write().await;
let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
// Get the provider first, before starting the reply stream
let provider = agent.provider().await;
// Create a single message for the prompt
let messages = vec![Message::user().with_text(request.prompt)];
// Get response from agent
let mut response_text = String::new();
let mut stream = match agent.reply(&messages).await {
let mut stream = match agent
.reply(
&messages,
Some(session::Identifier::Name(session_id.clone())),
)
.await
{
Ok(stream) => stream,
Err(e) => {
tracing::error!("Failed to start reply stream: {:?}", e);
@@ -254,15 +300,20 @@ async fn ask_handler(
}
};
// Collect all messages for storage
let mut all_messages = messages.clone();
let mut response_message = Message::assistant();
while let Some(response) = stream.next().await {
match response {
Ok(message) => {
if message.role == Role::Assistant {
for content in message.content {
for content in &message.content {
if let MessageContent::Text(text) = content {
response_text.push_str(&text.text);
response_text.push('\n');
}
response_message.content.push(content.clone());
}
}
}
@@ -273,6 +324,24 @@ async fn ask_handler(
}
}
// Add the complete response message to the conversation history
if !response_message.content.is_empty() {
all_messages.push(response_message);
}
// Get the session path - file will be created when needed
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
// Store messages and generate description in background
let session_path = session_path.clone();
let messages = all_messages.clone();
let provider = provider.clone();
tokio::spawn(async move {
if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
tracing::error!("Failed to store session history: {:?}", e);
}
});
Ok(Json(AskResponse {
response: response_text.trim().to_string(),
}))
@@ -394,6 +463,7 @@ mod tests {
.body(Body::from(
serde_json::to_string(&AskRequest {
prompt: "test prompt".to_string(),
session_id: Some("test-session".to_string()),
})
.unwrap(),
))
crates/goose-server/src/routes/session.rs
+128
View File
@@ -0,0 +1,128 @@
use crate::state::AppState;
use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
routing::get,
Json, Router,
};
use goose::message::Message;
use goose::session;
use serde::Serialize;
#[derive(Serialize)]
struct SessionInfo {
id: String,
path: String,
modified: String,
metadata: session::SessionMetadata,
}
#[derive(Serialize)]
struct SessionListResponse {
sessions: Vec<SessionInfo>,
}
#[derive(Serialize)]
struct SessionHistoryResponse {
session_id: String,
metadata: session::SessionMetadata,
messages: Vec<Message>,
}
// List all available sessions
async fn list_sessions(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<SessionListResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
let sessions = match session::list_sessions() {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to list sessions: {:?}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
let session_infos = sessions
.into_iter()
.map(|(id, path)| {
// Get last modified time as string
let modified = path
.metadata()
.and_then(|m| m.modified())
.map(|time| {
chrono::DateTime::<chrono::Utc>::from(time)
.format("%Y-%m-%d %H:%M:%S UTC")
.to_string()
})
.unwrap_or_else(|_| "Unknown".to_string());
// Get session description
let metadata = session::read_metadata(&path).expect("Failed to read session metadata");
SessionInfo {
id,
path: path.to_string_lossy().to_string(),
modified,
metadata,
}
})
.collect();
Ok(Json(SessionListResponse {
sessions: session_infos,
}))
}
// Get a specific session's history
async fn get_session_history(
State(state): State<AppState>,
headers: HeaderMap,
Path(session_id): Path<String>,
) -> Result<Json<SessionHistoryResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
// Read metadata
let metadata = session::read_metadata(&session_path).map_err(|_| StatusCode::NOT_FOUND)?;
let messages = match session::read_messages(&session_path) {
Ok(messages) => messages,
Err(e) => {
tracing::error!("Failed to read session messages: {:?}", e);
return Err(StatusCode::NOT_FOUND);
}
};
Ok(Json(SessionHistoryResponse {
session_id,
metadata,
messages,
}))
}
// Configure routes for this module
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/sessions", get(list_sessions))
.route("/sessions/:session_id", get(get_session_history))
.with_state(state)
}