stakpak_api/local/hooks/file_scratchpad_context/
mod.rs

1use crate::local::context_managers::file_scratchpad_context_manager::{
2    FileScratchpadContextManager, FileScratchpadContextManagerOptions,
3};
4use crate::models::AgentState;
5use stakpak_shared::define_hook;
6use stakpak_shared::hooks::{Hook, HookAction, HookContext, HookError, LifecycleEvent};
7use stakpak_shared::models::integrations::openai::Role;
8use stakpak_shared::models::llm::{LLMInput, LLMMessage, LLMMessageContent};
9
10const SYSTEM_PROMPT: &str = include_str!("./system_prompt.txt");
11const SCRATCHPAD_FILE: &str = ".stakpak/session/scratchpad.md";
12const TODO_FILE: &str = ".stakpak/session/todo.md";
13
14pub struct FileScratchpadContextHook {
15    pub context_manager: FileScratchpadContextManager,
16}
17
18pub struct FileScratchpadContextHookOptions {
19    pub scratchpad_path: Option<String>,
20    pub todo_path: Option<String>,
21    pub history_action_message_size_limit: Option<usize>,
22    pub history_action_message_keep_last_n: Option<usize>,
23    pub history_action_result_keep_last_n: Option<usize>,
24    pub overwrite_if_different: Option<bool>,
25}
26
27impl FileScratchpadContextHook {
28    pub fn new(options: FileScratchpadContextHookOptions) -> Self {
29        let context_manager =
30            FileScratchpadContextManager::new(FileScratchpadContextManagerOptions {
31                scratchpad_file_path: options
32                    .scratchpad_path
33                    .unwrap_or(SCRATCHPAD_FILE.to_string())
34                    .into(),
35                todo_file_path: options.todo_path.unwrap_or(TODO_FILE.to_string()).into(),
36                history_action_message_size_limit: options
37                    .history_action_message_size_limit
38                    .unwrap_or(100),
39                history_action_message_keep_last_n: options
40                    .history_action_message_keep_last_n
41                    .unwrap_or(1),
42                history_action_result_keep_last_n: options
43                    .history_action_result_keep_last_n
44                    .unwrap_or(50),
45                overwrite_if_different: options.overwrite_if_different.unwrap_or(true),
46            });
47
48        Self { context_manager }
49    }
50}
51
52define_hook!(
53    FileScratchpadContextHook,
54    "file_scratchpad_context",
55    async |&self, ctx: &mut HookContext<AgentState>, event: &LifecycleEvent| {
56        if *event != LifecycleEvent::BeforeInference {
57            return Ok(HookAction::Continue);
58        }
59
60        let model = ctx.state.active_model.clone();
61
62        let tools = ctx
63            .state
64            .tools
65            .clone()
66            .map(|t| t.into_iter().map(Into::into).collect());
67
68        let cwd = std::env::current_dir().unwrap_or_default();
69        let scratchpad_file_path =
70            cwd.join(self.context_manager.get_scratchpad_path(ctx.session_id));
71        let todo_file_path = cwd.join(self.context_manager.get_todo_path(ctx.session_id));
72        let system_prompt = SYSTEM_PROMPT
73            .replace(
74                "{{SCRATCHPAD_PATH}}",
75                &scratchpad_file_path.display().to_string(),
76            )
77            .replace("{{TODO_PATH}}", &todo_file_path.display().to_string());
78
79        let mut messages = Vec::new();
80        messages.push(LLMMessage {
81            role: Role::System.to_string(),
82            content: LLMMessageContent::String(system_prompt),
83        });
84        messages.extend(
85            self.context_manager
86                .reduce_context_with_session(ctx.state.messages.clone(), ctx.session_id),
87        );
88
89        ctx.state.llm_input = Some(LLMInput {
90            model,
91            messages,
92            max_tokens: 16000,
93            tools,
94            provider_options: None,
95            headers: None,
96        });
97
98        Ok(HookAction::Continue)
99    }
100);
stakpak_api/local/hooks/file_scratchpad_context/mod.rs

stakpak_api/local/hooks/file_scratchpad_context/
mod.rs
