Add context dump/load feature

Adds session context management to prevent session poisoning:
- CONTEXT_DIR and ENABLE_CONTEXT_DUMP config options
- issue_ref_to_context_file() - derive context file path
- kugetsu_context_load() - load previous context
- kugetsu_context_dump() - save context on session start
- kugetsu_context_update_message() - append to conversation history
- Integration in cmd_start and cmd_continue
- New 'kugetsu context' command
This commit is contained in:
shokollm
2026-04-05 04:23:26 +00:00
parent b76a9b883a
commit 61f06f825f

View File

@@ -13,6 +13,8 @@ VERBOSITY_DIR="$KUGETSU_DIR/verbosity"
MAX_CONCURRENT_AGENTS="${MAX_CONCURRENT_AGENTS:-3}"
KUGETSU_VERBOSITY="${KUGETSU_VERBOSITY:-default}"
CONTEXT_DIR="${CONTEXT_DIR:-$KUGETSU_DIR/context}"
ENABLE_CONTEXT_DUMP="${ENABLE_CONTEXT_DUMP:-true}"
WORKTREE_CHECK_PR_STATUS="${WORKTREE_CHECK_PR_STATUS:-true}"
# Load user config overrides (~/.kugetsu/config)
@@ -79,6 +81,7 @@ Usage:
kugetsu destroy --pm-agent [-y] Delete pm-agent session (not recommended)
kugetsu destroy --base [-y] Delete base session
kugetsu set-pr <issue-ref> <pr-url> Set PR URL for session (for PR tracking)
kugetsu context <issue-ref> Show context for issue
kugetsu help Show this help
Issue Ref Format:
@@ -317,6 +320,163 @@ filename_to_issue_ref() {
echo "$name" | sed 's/-\([0-9]*\)$/#\1' | sed 's/-/\//g'
}
issue_ref_to_context_file() {
local issue_ref="$1"
local context_filename=$(issue_ref_to_filename "$issue_ref")
echo "$CONTEXT_DIR/${context_filename}.json"
}
kugetsu_context_load() {
local issue_ref="$1"
if [ "$ENABLE_CONTEXT_DUMP" != "true" ]; then
echo ""
return
fi
local context_file=$(issue_ref_to_context_file "$issue_ref")
if [ ! -f "$context_file" ]; then
echo ""
return
fi
python3 << PYEOF
import json
import sys
context_file = "$context_file"
try:
with open(context_file, 'r') as f:
ctx = json.load(f)
lines = []
lines.append("## PREVIOUS CONTEXT")
lines.append(f"Issue: {ctx.get('issue_ref', 'unknown')}")
lines.append(f"Last updated: {ctx.get('updated_at', 'unknown')}")
lines.append(f"Current branch: {ctx.get('current_branch', 'unknown')}")
lines.append("")
lines.append("### Previous work summary:")
lines.append(ctx.get('last_message', '(no previous message)'))
lines.append("")
history = ctx.get('conversation_history', [])
if history:
lines.append("### Conversation history:")
for msg in history[-5:]:
role = msg.get('role', 'unknown')
content = msg.get('content', '')
ts = msg.get('timestamp', '')
lines.append(f"- [{ts}] {role}: {content[:200]}...")
print('\n'.join(lines))
except Exception as e:
print(f"Warning: Failed to load context: {e}", file=sys.stderr)
print("", file=sys.stderr)
PYEOF
}
kugetsu_context_dump() {
local issue_ref="$1"
local message="$2"
local branch_name="${3:-}"
if [ "$ENABLE_CONTEXT_DUMP" != "true" ]; then
return
fi
local context_file=$(issue_ref_to_context_file "$issue_ref")
mkdir -p "$CONTEXT_DIR"
python3 << PYEOF
import json
import os
from datetime import datetime
context_file = "$context_file"
issue_ref = "$issue_ref"
message = """$message"""
branch_name = "$branch_name"
context = {
"issue_ref": issue_ref,
"current_branch": branch_name,
"updated_at": datetime.now().isoformat() + "Z",
"last_message": message[:500] if message else "",
"conversation_history": []
}
if os.path.exists(context_file):
try:
with open(context_file, 'r') as f:
existing = json.load(f)
history = existing.get('conversation_history', [])
history.append({
"role": "user",
"content": message[:1000] if message else "",
"timestamp": datetime.now().isoformat() + "Z"
})
history = history[-20:]
context["conversation_history"] = history
context["created_at"] = existing.get("created_at", context["updated_at"])
except:
context["created_at"] = datetime.now().isoformat() + "Z"
else:
context["created_at"] = datetime.now().isoformat() + "Z"
with open(context_file, 'w') as f:
json.dump(context, f, indent=2)
PYEOF
}
kugetsu_context_update_message() {
local issue_ref="$1"
local message="$2"
if [ "$ENABLE_CONTEXT_DUMP" != "true" ]; then
return
fi
local context_file=$(issue_ref_to_context_file "$issue_ref")
if [ ! -f "$context_file" ]; then
return
fi
python3 << PYEOF
import json
from datetime import datetime
context_file = "$context_file"
message = """$message"""
try:
with open(context_file, 'r') as f:
ctx = json.load(f)
history = ctx.get('conversation_history', [])
history.append({
"role": "assistant",
"content": message[:1000] if message else "",
"timestamp": datetime.now().isoformat() + "Z"
})
history = history[-20:]
ctx["conversation_history"] = history
ctx["last_message"] = message[:500] if message else ""
ctx["updated_at"] = datetime.now().isoformat() + "Z"
with open(context_file, 'w') as f:
json.dump(ctx, f, indent=2)
except Exception as e:
pass
PYEOF
}
update_session_pr_url() {
local issue_ref="$1"
local pr_url="$2"
@@ -1557,7 +1717,10 @@ cmd_start() {
> "$fork_log"
local fork_context=$(kugetsu_get_fork_context "$issue_ref")
local previous_context=$(kugetsu_context_load "$issue_ref")
local branch_name=$(issue_ref_to_branch_name "$issue_ref")
local full_message="${fork_context}
${previous_context}
## YOUR TASK
$message"
@@ -1659,6 +1822,8 @@ PYEOF
add_issue_to_index "$issue_ref" "$session_file"
kugetsu_context_dump "$issue_ref" "$message" "$branch_name"
echo "Session started for '$issue_ref': $new_session_id"
echo "Worktree: $worktree_path"
}
@@ -1706,22 +1871,31 @@ cmd_continue() {
local worktree_path=$(python3 -c "import json; print(json.load(open('$session_path')).get('worktree_path', ''))" 2>/dev/null || echo "")
echo "Continuing session for '$session_name'..."
local previous_context=$(kugetsu_context_load "$session_name")
local full_message="${previous_context}
## CONTINUE TASK
$message"
# Note: --continue always allowed (existing sessions don't count toward limit)
# Wrap in subshell with cd to ensure worktree directory is set correctly in session DB
if [ -n "$worktree_path" ] && [ -d "$worktree_path" ]; then
echo "Using worktree: $worktree_path"
if [ "$DEBUG_MODE" = true ]; then
(cd "$worktree_path" && opencode run "$message" --continue --session "$opencode_session_id" --dir "$worktree_path" 2>&1) | tee "$session_path.debug.log" &
(cd "$worktree_path" && opencode run "$full_message" --continue --session "$opencode_session_id" --dir "$worktree_path" 2>&1) | tee "$session_path.debug.log" &
else
(cd "$worktree_path" && opencode run "$message" --continue --session "$opencode_session_id" --dir "$worktree_path" 2>&1) &
(cd "$worktree_path" && opencode run "$full_message" --continue --session "$opencode_session_id" --dir "$worktree_path" 2>&1) &
fi
else
if [ "$DEBUG_MODE" = true ]; then
opencode run "$message" --continue --session "$opencode_session_id" 2>&1 | tee "$session_path.debug.log" &
opencode run "$full_message" --continue --session "$opencode_session_id" 2>&1 | tee "$session_path.debug.log" &
else
opencode run "$message" --continue --session "$opencode_session_id" 2>&1 &
opencode run "$full_message" --continue --session "$opencode_session_id" 2>&1 &
fi
fi
kugetsu_context_update_message "$session_name" "$message"
}
cmd_list() {
@@ -2028,6 +2202,22 @@ main() {
validate_issue_ref "$issue_ref"
update_session_pr_url "$issue_ref" "$pr_url"
;;
context)
local issue_ref="${1:-}"
if [ -z "$issue_ref" ]; then
echo "Usage: kugetsu context <issue-ref>" >&2
echo "Show context for an issue" >&2
exit 1
fi
validate_issue_ref "$issue_ref"
local context_file=$(issue_ref_to_context_file "$issue_ref")
if [ -f "$context_file" ]; then
cat "$context_file"
else
echo "No context found for '$issue_ref'" >&2
exit 1
fi
;;
*)
echo "Error: unknown command '$command'" >&2
usage