@@ -40,11 +40,13 @@ use crate::skill::SkillRegistry;
4040use crate :: tool:: { Registry , ToolContext , ToolExecutionMode } ;
4141use anyhow:: Result ;
4242use futures:: StreamExt ;
43+ use std:: cell:: Cell ;
4344use std:: collections:: { HashMap , HashSet } ;
4445use std:: hash:: { Hash , Hasher } ;
4546use std:: io:: { self , Write } ;
4647use std:: path:: PathBuf ;
47- use std:: sync:: { Arc , LazyLock , Mutex as StdMutex } ;
48+ use std:: ptr:: NonNull ;
49+ use std:: sync:: { Arc , LazyLock , Mutex as StdMutex , Mutex } ;
4850use std:: time:: { Duration , Instant } ;
4951use tokio:: sync:: { broadcast, mpsc} ;
5052
@@ -69,6 +71,40 @@ static JCODE_REPO_SOURCE_STATE: LazyLock<(Option<String>, Option<bool>)> = LazyL
6971} ) ;
7072static WORKING_GIT_STATE_CACHE : LazyLock < StdMutex < HashMap < PathBuf , Option < GitState > > > > =
7173 LazyLock :: new ( || StdMutex :: new ( HashMap :: new ( ) ) ) ;
74+
75+ // Thread-local agent accessor for DCP tools to reach the agent's DCP plugin
76+ #[ cfg( feature = "dcp" ) ]
77+ thread_local ! {
78+ static CURRENT_AGENT : Cell <Option <NonNull <Agent >>> = const { Cell :: new( None ) } ;
79+ }
80+
81+ #[ cfg( feature = "dcp" ) ]
82+ impl Agent {
83+ /// Set the current agent pointer (called by the server before tool execution)
84+ pub ( crate ) fn set_current_agent_ptr ( & mut self ) {
85+ CURRENT_AGENT . with ( |cell| {
86+ cell. set ( Some ( NonNull :: from ( self ) ) ) ;
87+ } ) ;
88+ }
89+
90+ /// Clear the current agent pointer (called when agent is dropped)
91+ pub ( crate ) fn clear_current_agent_ptr ( ) {
92+ CURRENT_AGENT . with ( |cell| {
93+ cell. set ( None ) ;
94+ } ) ;
95+ }
96+
97+ /// Access the current agent's DCP plugin for tool execution
98+ pub ( crate ) fn get_current_dcp ( ) -> Option < Arc < Mutex < crate :: dcp_plugin:: DcpPlugin > > > {
99+ CURRENT_AGENT . with ( |cell| {
100+ cell. get ( ) . map ( |ptr| {
101+ // SAFETY: ptr is guaranteed valid for the duration of tool execution
102+ // because the server clears it when the agent is dropped
103+ unsafe { ptr. as_ref ( ) } . registry . dcp ( )
104+ } ) . flatten ( )
105+ } )
106+ }
107+ }
72108const STREAM_KEEPALIVE_PONG_ID : u64 = 0 ;
73109
74110fn stable_hash_str ( value : & str ) -> u64 {
@@ -321,6 +357,13 @@ impl Agent {
321357 agent. session . provider_key =
322358 crate :: session:: derive_session_provider_key ( agent. provider . name ( ) ) ;
323359 agent. session . ensure_initial_session_context_message ( ) ;
360+
361+ // Wire DCP plugin into registry so DCP tools can access it
362+ #[ cfg( feature = "dcp" ) ]
363+ if let Some ( dcp) = agent. dcp . take ( ) {
364+ agent. registry . set_dcp ( dcp) ;
365+ }
366+
324367 agent. seed_compaction_from_session ( ) ;
325368 agent. log_env_snapshot ( "create" ) ;
326369 crate :: telemetry:: begin_session_with_parent (
@@ -373,6 +416,13 @@ impl Agent {
373416 agent. restore_reasoning_effort_from_session ( ) ;
374417 agent. session . ensure_initial_session_context_message ( ) ;
375418 agent. sync_memory_dedup_state_from_session ( ) ;
419+
420+ // Wire DCP plugin into registry so DCP tools can access it
421+ #[ cfg( feature = "dcp" ) ]
422+ if let Some ( dcp) = agent. dcp . take ( ) {
423+ agent. registry . set_dcp ( dcp) ;
424+ }
425+
376426 agent. seed_compaction_from_session ( ) ;
377427 agent. log_env_snapshot ( "attach" ) ;
378428 crate :: telemetry:: begin_session_with_parent (
@@ -417,6 +467,13 @@ impl Agent {
417467 self . session. messages. len( )
418468 ) ) ;
419469 drop ( manager) ;
470+
471+ // J9: Set session ID on DCP pruner for proper persistence
472+ #[ cfg( feature = "dcp" ) ]
473+ if let Some ( ref mut dcp) = self . dcp {
474+ dcp. pruner_mut ( ) . set_session_id ( & self . session . id ) ;
475+ }
476+
420477 if let Some ( state) = sanitized_state {
421478 self . session . compaction = state;
422479 self . persist_session_best_effort ( "sanitized oversized OpenAI native compaction" ) ;
0 commit comments