Fix: clean LLM configuration

[calm fulcrum]

🧵 actually I am stuck on something stil... 😭

FYI @cerulean zinc @brazen burrow @remote laurel

I now have a way to persist the LLM config session-wide, and share it across all clients in the session...

[remote laurel]

Awesome 👏

[calm fulcrum]

But I'm missing a way to query the main client in the session, when I'm currently procesisng a query from another client

[calm fulcrum]

Still lost in the maze of engine code 😭

@reef cobalt any chance I could bother you for a few minutes? 🙏

[reef cobalt]

I'm getting release stuff ready today but happy to help async while waiting for things to run 🙂

[calm fulcrum]

Oh yeah async is fine, if I can manage to formulate my question in writing, I think you will find it easy to answer 😛

• Basically I'm deep within a dagql query in core/llm.go so I have access to my client's dagql.Server, core.Query, and from there I can get the Server interface
• What I want to do is call dagql.Server.Select() on another client, specifically the main client of my current client's session.
• The closest I got is Server.MainClientID(). But then how to get a full working dagql.Server instance from there, I don't know

Also, maybe I don't even need to get a dagql.Server. I'm only using that to select secret("env://OPENAI_API_KEY") and other client-side secrets that I need to get from the main client of my session. Any method for getting that info would be fine (access the client's secret store directly? couldn't navigate that API though)

If it's easier, I can also do it from the engine/server side, for example in getOrInitDaggerClient

diff --git a/core/llm.go b/core/llm.go
index 80f05cd1a..19498de15 100644
--- a/core/llm.go
+++ b/core/llm.go
@@ -251,11 +251,25 @@ func loadGlobalLlmRouter(ctx context.Context, srv *dagql.Server) (*LlmRouter, er
 }
 
 func NewLlm(ctx context.Context, query *Query, srv *dagql.Server, model string) (*Llm, error) {
-    // FIXME: finish dismantling the global llm config machinery
-    router, err := loadGlobalLlmRouter(ctx, srv)
+    routerData, err := query.Server.GetUserData(ctx, "llm-router")
     if err != nil {
-        return nil, err
+        return nil.err
+    }
+    if routerData == nil {
+        mainClientID, err := query.Server.MainClientCallerID(ctx)
+        if err != nil {
+            return nil, err
+        }
+        // FIXME: how to get secrets out of the main client from here?
+        routerData, err = NO_IDEA_HOW_TO_DO_THIS(ctx, mainClientID)
+        if err != nil {
+            return nil, err
+        }
+        if err := query.Server.SetUserData(ctx, "llm-router", router); err != nil {
+            return nil, err
+        }
     }
+    router := routerData.(*LlmRouter)
     if model == "" {
         model = router.DefaultModel()
     }
diff --git a/core/query.go b/core/query.go
index 25781e78c..a0386774e 100644
--- a/core/query.go
+++ b/core/query.go
@@ -106,6 +106,12 @@ type Server interface {
 
     // Open a client's SQLite database. Be sure to close it!
     OpenClientDB(clientID string) (*sql.DB, error)
+
+    // Retrieve session-wide user data
+    GetUserData(context.Context, string) (any, error)
+
+    // Attach session-wide user data
+    SetUserData(context.Context, string, any) error
 }
 
 func NewRoot(srv Server) *Query {
diff --git a/engine/server/session.go b/engine/server/session.go
index a321f0edc..9f1e6188e 100644
--- a/engine/server/session.go
+++ b/engine/server/session.go
@@ -92,6 +92,9 @@ type daggerSession struct {
 
     interactive        bool
     interactiveCommand []string
+
+    userdata     map[string]any
+    userdataLock sync.RWMutex
 }
 
 type daggerSessionState string
@@ -1265,6 +1268,37 @@ func (srv *Server) MainClientCallerID(ctx context.Context) (string, error) {
     return client.daggerSession.mainClientCallerID, nil
 }
 
+// GetUserData retrieves session-wide user data by key.
+// If the user data is not set, nil is returned
+func (srv *Server) GetUserData(ctx context.Context, key string) (any, error) {
+    client, err := srv.clientFromContext(ctx)
+    if err != nil {
+        return nil, err
+    }
+    client.daggerSession.userdataLock.RLock()
+    defer client.daggerSession.userdataLock.RUnlock()
+    data, ok := client.daggerSession.userdata[key]
+    if !ok {
+        return nil, nil
+    }
+    return data, nil
+}
+
+// SetUserData attaches session-wide user data by key
+func (srv *Server) SetUserData(ctx context.Context, key string, data any) error {
+    client, err := srv.clientFromContext(ctx)
+    if err != nil {
+        return err
+    }
+    client.daggerSession.userdataLock.Lock()
+    defer client.daggerSession.userdataLock.Unlock()
+    if client.daggerSession.userdata == nil {
+        client.daggerSession.userdata = make(map[string]any)
+    }
+    client.daggerSession.userdata[key] = data
+    return nil
+}
+
 // The default deps of every user module (currently just core)
 func (srv *Server) DefaultDeps(ctx context.Context) (*core.ModDeps, error) {
     client, err := srv.clientFromContext(ctx)

[reef cobalt]

Okay that makes sense, I think just adding a method to Query/Server that lets you make a dagql.Select in a given client's server would be fine if nothing else.

But to double check something, do you specifically need an instance of a core.Secret? Or do you just need the plaintext of the secret? If you just need the plaintext there might be an even easier option

[calm fulcrum]

I really just need the plaintexts (plural)

Even the userdata idea is optional if there's a better way. It seemed to make sense to use it as a "cache", so once hydrated, it's accessible to all clients in the session

I could even combine it all into:

type Server interface {
  // Retrieve the secret from the session's main client
  // Each secret will only be retrieved once, then cached for the session
  GetSessionSecret(uri string) (string, error)
}

Then I can tuck the "secret retrieval cache" in the session. And that will take care of the "cached hydration" logic

(at the moment I don't have any other user data I need to attach, it's only for session-wide sharing of main client's secrets relative to llm config)

There's a slightly annoying side effect to this method, which is that all the secret retrievals are visible in the trace - so when some of the env variables are not set, you see scary errors. But I'll worry about that later.

[reef cobalt]

Okay if you just need plaintexts then you could probably just add a slight variant of this existing method Server.Secrets: https://github.com/sipsma/dagger/blob/f8fdc7b88a4e66e6d0c60f719321cfef180d915a/engine/server/session.go#L1312-L1312

That one returns the secret store for the current client, but you could just add a separate method like MainClientSecrets that returns the main client's secret store.

Looking at the current methods on SecretStore, they all expect the id digest for the secret, but I suppose you could just add another quick one that lets you ask for the secret with a given URI (just iterate over the secrets and find the match): https://github.com/sipsma/dagger/blob/f8fdc7b88a4e66e6d0c60f719321cfef180d915a/core/secret.go#L48-L48

I guess the only downside to that is what you hinted at in that the secrets will always be retrieved on demand from the client. But they won't be visible in the trace. I think Kyle added some basic client-side caching of secret providers too so you at least won't be reinvoking vault/op/etc. every time you retrieve the secret, so I wouldn't worry too much about the overhead until proven wrong

[calm fulcrum]

Oh yeah great. If I can skip implementing my own caching, even better

thanks I'll try that!

@reef cobalt so just to make sure I got it:

I could implement eg. func (s *SecretStore) LookupPlaintext(uri string) (string, error), that would:

• First instantiate a new core.Secret for that URI (not sure how?)
• Then add it to the main client's secret store with SecrrtStore.AddSecret()
• Then iterate over the secret store and find the right digest?

OK I'm confused as to how to implement this, sorry

[reef cobalt]

Oh wait, okay I overlooked the fact that the main client never actually explicitly created the secret in the first place. What I explained only makes sense if the secret you want is already loaded by the client, sorry, this is slightly different than what we've done w/ secrets in the past.

Hm in that case actually yeah it's probably easiest for now to just do a dagql.Select to get the secret plaintext. It shouldn't be too hard to add a method to Server.

The only adjustment needed is that right now the dagql.Server for each client is just a local variable, but you want to save it on the daggerClient struct. So just make this dag var a field that's set on the client: https://github.com/sipsma/dagger/blob/68f8c901f2c5666f8e182c472f6cf9ab548fc713/engine/server/session.go#L558-L558

Then you'd just need a method on Server along the lines of

func (srv *Server) LoadSecretPlaintextFromMainClient(ctx context.Context, uri string) (string, error) {
    clientMetadata, err := engine.ClientMetadataFromContext(ctx)
    if err != nil {
        // ...
    }
    mainClient, ok := srv.clientFromIDs(clientMetadata.SessionID, client.daggerSession.mainClientCallerID)
    if !ok {
        // ...
    }
    err = mainClient.dag.Select(...) // select the secret plaintext for the given uri and then return it
}

Or anything like that. If you think you need more variations on it you could make the method more generic like func (srv *Server) SelectFromMainClient(ctx context.Context, root, object dagql.Object, selectors ...dagql.Selector), basically just the same signature as dagql.Select but implemented to run against the main client

[calm fulcrum]

Ok I see, thanks. I was trying to follow how dag was being used, and how I might get it back. Now the answer is clear 🙂

@reef cobalt do you think this is enough?

@@ -553,6 +556,7 @@ func (srv *Server) initializeDaggerClient(
        if err := coreMod.Install(ctx, dag); err != nil {
                return fmt.Errorf("failed to install core module: %w", err)
        }
+       client.dag = dag
        client.defaultDeps = core.NewModDeps(client.dagqlRoot, []core.Mod{coreMod})

        if opts.EncodedModuleID == "" {

[reef cobalt]

yep that should be it

[calm fulcrum]

Awesome, I think I have what I need, trying now. Thank you!

Mmmm @reef cobalt how do I get the main client object from another client in the session? The closest I can get is client.daggerSession.mainClientCallerID, is there an easy way to get a daggerClient instance from that ID?

Oh I see a clients map[string]*daggerClient... guessing thats it

[reef cobalt]

that's what the first 8 lines of the snippet here does #1341485480591626262 message

[calm fulcrum]

oh right, sorry!

https://tenor.com/view/hose-gif-20117424

my brain rn

[calm fulcrum]

Looks like it worked!

Mmm looks like not completely working yet

@reef cobalt so my implementation half-works:

• It correctly gets a dagql.Server from the client object, all the plumbing is in place as we discussed it
• But the dagql.Server I retrieve from a module client does not seem to be the one for the main client. It seems to behave exactly like the current client, even though I retrieved it via clients[mainclientID].

So everything works when calling from the CLI. But does not work when called from a module

confirmed that it's the current client

https://github.com/shykes/dagger/commit/5ea3dd538a66a2bebb0d3651350ba0d07e38782e

[reef cobalt]

Oh, you might need to update the ctx to have the right metadata before providing it to the dagql.Select call. We stuff the current client metadata in there and use it in various places across the api. Try

currentClientMetadata, err := engine.ClientMetadataFromContext(ctx)
if err != nil { ... }
mainClient, ok := srv.clientFromIDs(currentClientMetadata.SessionID, currentClient.daggerSession.mainClientCallerID)
if !ok { ... }

ctx = engine.ContextWithClientMetadata(ctx, mainClient.clientMetadata)
err = mainClient.dag.Select(ctx, ...)

Oh I just saw the commit you posted, I guess you were going a slightly different direction, but same idea either way. Need to put the main client metadata in the ctx that's provided to dag.Select

[calm fulcrum]

Is this something I would do in my Server.MainClient() implementation? Or on the caller side in core?

[reef cobalt]

I feel like it's easier to setup the context in the server pkg just because it has direct access to the .clientMetadata field already. So you could update your (srv *Server) MainServer method to also return a context.Context that's setup with the metadata already

Or anything equivalent to that

[calm fulcrum]

return a context?

ok similar pattern to eg. our tracer

(note I went for MainServer instead of eg. GetMainSecret because I already had all the logic on the caller side to use a dagql.Server, made it easier to swap in that way)

New implementation:

func (srv *Server) MainServer(ctx context.Context) (context.Context, *dagql.Server, error) {
    client, err := srv.clientFromContext(ctx)
    if err != nil {
        return nil, nil, err
    }
    clientMetadata, err := engine.ClientMetadataFromContext(ctx)
    if err != nil {
        return nil, nil, err
    }
    mainClient, ok := srv.clientFromIDs(clientMetadata.SessionID, client.daggerSession.mainClientCallerID)
    if !ok {
        return nil, nil, fmt.Errorf("failed to retrieve session main client")
    }
    ctx = engine.ContextWithClientMetadata(ctx, mainClient.clientMetadata)
    return ctx, mainClient.dag, nil
}

Is it important to only use that returned context for the secret lookup, and only that?

[reef cobalt]

Yes I believe so. I'm obviously going off theory here but not using a context w/ the main client's metadata sounds like it would result in exactly the problem you were seeing, so good chance this will fix it. If not I'll take a closer look

[calm fulcrum]

testing now!

it worked! thank you so much @reef cobalt